/* C++11 implementation for "Selection in X + Y and matrices with sorted
rows and columns" by A. Mirzaian and E. Arjomandi.
http://w...content-available-to-author-only...u.ca/~andy/pubs/X%2BY.pdf
*/
#include <algorithm>
#include <vector>
#include <iostream>
#include <assert.h>
int qceil(int x)
{
return (x + 3) / 4;
}
int qhigh(int x, int n)
{
if (n & 1) // odd
return qceil(x + 2 * n + 1);
else // even
return n + 1 + qceil(x);
}
typedef std::vector<int> Vec;
struct Ranks
{
int ra_less;
int rb_more;
Vec l;
Ranks(const Vec& a1, const Vec& a2, int a, int b)
{
int n = a1.size() - 1;
ra_less = 0;
rb_more = n * n;
l.reserve(12 * n + 1);
l.push_back(0);
int j = n;
for (int i = 1; i <= n; ++i)
{
while (j && a1[i] + a2[j] <= a)
--j;
ra_less += j;
int jj = j;
while (jj && a1[i] + a2[jj] < b)
{
l.push_back(a1[i] + a2[jj]);
--jj;
}
rb_more -= jj;
}
}
};
Vec half(const Vec& a)
{
Vec res;
res.reserve(2 + (a.size() - 1) / 2);
res.push_back(0);
for (int i = 1; i < a.size(); i += 2)
res.push_back(a[i]);
if (a.size() & 1)
res.push_back(a[a.size() - 1]);
return res;
}
struct AB
{
int a;
int b;
};
AB biselect(const Vec& a1, const Vec& a2, int k1, int k2)
{
AB res;
int n = a1.size() - 1;
assert(n > 0);
if (n == 1)
{
res.a = res.b = a1[1] + a2[1];
}
else if (n == 2)
{
Vec l {a1[1] + a2[1], a1[1] + a2[2], a1[2] + a2[1], a1[2] + a2[2]};
std::nth_element(l.begin(),
l.end() - k1,
l.end());
res.a = l[4 - k1];
std::nth_element(l.begin(),
l.end() - k2,
l.end());
res.b = l[4 - k2];
}
else
{
int k1_half = qhigh(k1, n);
int k2_half = qceil(k2);
AB ab = biselect(half(a1), half(a2), k1_half, k2_half);
Ranks ranks(a1, a2, ab.a, ab.b);
int r1 = k1 + ranks.rb_more - n * n;
int r2 = k2 + ranks.rb_more - n * n;
if (ranks.ra_less <= k1 - 1)
res.a = ab.a;
else if (r1 <= 0)
res.a = ab.b;
else
{
std::nth_element(ranks.l.begin() + 1,
ranks.l.end() - r1,
ranks.l.end());
res.a = *(ranks.l.end() - r1);
}
if (ranks.ra_less <= k2 - 1)
res.b = ab.a;
else if (r2 <= 0)
res.b = ab.b;
else
{
std::nth_element(ranks.l.begin() + 1,
ranks.l.end() - r2,
ranks.l.end());
res.b = *(ranks.l.end() - r2);
}
}
return res;
}
int select(const Vec& a1, const Vec& a2, int k)
{
assert(a1.size() == a2.size());
AB ab = biselect(a1, a2, k, k);
assert(ab.a == ab.b);
return ab.a;
}
int check(const Vec& a1, const Vec& a2, int k)
{
int n = a1.size() - 1;
Vec sum;
sum.reserve(n * n);
for (int i = 1; i <= n; ++i)
for (int j = 1; j <= n; ++j)
sum.push_back(a1[i] + a2[j]);
std::nth_element(sum.begin(),
sum.end() - k,
sum.end());
return *(sum.end() - k);
}
#include <random>
int main()
{
std::random_device rd;
std::default_random_engine e1(rd());
std::uniform_int_distribution<int> uniform_dist(1, 6);
int n = 100000;
Vec a1, a2;
a1.reserve(n+1);
a2.reserve(n+1);
int x = 1000;
a1.push_back(0);
for (int i = 0; i < n; ++i)
{
x -= uniform_dist(e1);
a1.push_back(x);
}
x = 1000;
a2.push_back(0);
for (int i = 0; i < n; ++i)
{
x -= uniform_dist(e1);
a2.push_back(x);
}
std::cout << select(a1, a2, 1 + n * n / 7) << '\n'
/*<< check(a1, a2, 1 + n * n / 7) << '\n'*/;
return 0;
}
LyogQysrMTEgaW1wbGVtZW50YXRpb24gZm9yICJTZWxlY3Rpb24gaW4gWCArIFkgYW5kIG1hdHJpY2VzIHdpdGggc29ydGVkCiAgIHJvd3MgYW5kIGNvbHVtbnMiIGJ5IEEuIE1pcnphaWFuIGFuZCBFLiBBcmpvbWFuZGkuCiAgIGh0dHA6Ly93Li4uY29udGVudC1hdmFpbGFibGUtdG8tYXV0aG9yLW9ubHkuLi51LmNhL35hbmR5L3B1YnMvWCUyQlkucGRmCiAqLwoKI2luY2x1ZGUgPGFsZ29yaXRobT4KI2luY2x1ZGUgPHZlY3Rvcj4KI2luY2x1ZGUgPGlvc3RyZWFtPgojaW5jbHVkZSA8YXNzZXJ0Lmg+CgppbnQgcWNlaWwoaW50IHgpCnsKICByZXR1cm4gKHggKyAzKSAvIDQ7Cn0KCmludCBxaGlnaChpbnQgeCwgaW50IG4pCnsKICBpZiAobiAmIDEpIC8vIG9kZAogICAgcmV0dXJuIHFjZWlsKHggKyAyICogbiArIDEpOwogIGVsc2UgLy8gZXZlbgogICAgcmV0dXJuIG4gKyAxICsgcWNlaWwoeCk7Cn0KCnR5cGVkZWYgc3RkOjp2ZWN0b3I8aW50PiBWZWM7CgpzdHJ1Y3QgUmFua3MKewogIGludCByYV9sZXNzOwogIGludCByYl9tb3JlOwogIFZlYyBsOwoKICBSYW5rcyhjb25zdCBWZWMmIGExLCBjb25zdCBWZWMmIGEyLCBpbnQgYSwgaW50IGIpCiAgewogICAgaW50IG4gPSBhMS5zaXplKCkgLSAxOwogICAgcmFfbGVzcyA9IDA7CiAgICByYl9tb3JlID0gbiAqIG47CiAgICBsLnJlc2VydmUoMTIgKiBuICsgMSk7CiAgICBsLnB1c2hfYmFjaygwKTsKCiAgICBpbnQgaiA9IG47CiAgICBmb3IgKGludCBpID0gMTsgaSA8PSBuOyArK2kpCiAgICB7CiAgICAgIHdoaWxlIChqICYmIGExW2ldICsgYTJbal0gPD0gYSkKCS0tajsKCiAgICAgIHJhX2xlc3MgKz0gajsKICAgICAgaW50IGpqID0gajsKCiAgICAgIHdoaWxlIChqaiAmJiBhMVtpXSArIGEyW2pqXSA8IGIpCiAgICAgIHsKCWwucHVzaF9iYWNrKGExW2ldICsgYTJbampdKTsKCS0tamo7CiAgICAgIH0KCiAgICAgIHJiX21vcmUgLT0gamo7CiAgICB9CiAgfQp9OwoKVmVjIGhhbGYoY29uc3QgVmVjJiBhKQp7CiAgVmVjIHJlczsKICByZXMucmVzZXJ2ZSgyICsgKGEuc2l6ZSgpIC0gMSkgLyAyKTsKICByZXMucHVzaF9iYWNrKDApOwoKICBmb3IgKGludCBpID0gMTsgaSA8IGEuc2l6ZSgpOyBpICs9IDIpCiAgICByZXMucHVzaF9iYWNrKGFbaV0pOwoKICBpZiAoYS5zaXplKCkgJiAxKQogICAgcmVzLnB1c2hfYmFjayhhW2Euc2l6ZSgpIC0gMV0pOwoKICByZXR1cm4gcmVzOwp9CgpzdHJ1Y3QgQUIKewogIGludCBhOwogIGludCBiOwp9OwoKQUIgYmlzZWxlY3QoY29uc3QgVmVjJiBhMSwgY29uc3QgVmVjJiBhMiwgaW50IGsxLCBpbnQgazIpCnsKICBBQiByZXM7CiAgaW50IG4gPSBhMS5zaXplKCkgLSAxOwogIGFzc2VydChuID4gMCk7CgogIGlmIChuID09IDEpCiAgewogICAgcmVzLmEgPSByZXMuYiA9IGExWzFdICsgYTJbMV07CiAgfQogIGVsc2UgaWYgKG4gPT0gMikKICB7CiAgICBWZWMgbCB7YTFbMV0gKyBhMlsxXSwgYTFbMV0gKyBhMlsyXSwgYTFbMl0gKyBhMlsxXSwgYTFbMl0gKyBhMlsyXX07CiAgICBzdGQ6Om50aF9lbGVtZW50KGwuYmVnaW4oKSwKCQkgICAgIGwuZW5kKCkgLSBrMSwKCQkgICAgIGwuZW5kKCkpOwogICAgcmVzLmEgPSBsWzQgLSBrMV07CiAgICBzdGQ6Om50aF9lbGVtZW50KGwuYmVnaW4oKSwKCQkgICAgIGwuZW5kKCkgLSBrMiwKCQkgICAgIGwuZW5kKCkpOwogICAgcmVzLmIgPSBsWzQgLSBrMl07CiAgfQogIGVsc2UKICB7CiAgICBpbnQgazFfaGFsZiA9IHFoaWdoKGsxLCBuKTsKICAgIGludCBrMl9oYWxmID0gcWNlaWwoazIpOwogICAgQUIgYWIgPSBiaXNlbGVjdChoYWxmKGExKSwgaGFsZihhMiksIGsxX2hhbGYsIGsyX2hhbGYpOwogICAgUmFua3MgcmFua3MoYTEsIGEyLCBhYi5hLCBhYi5iKTsKICAgIGludCByMSA9IGsxICsgcmFua3MucmJfbW9yZSAtIG4gKiBuOwogICAgaW50IHIyID0gazIgKyByYW5rcy5yYl9tb3JlIC0gbiAqIG47CgogICAgaWYgKHJhbmtzLnJhX2xlc3MgPD0gazEgLSAxKQogICAgICByZXMuYSA9IGFiLmE7CiAgICBlbHNlIGlmIChyMSA8PSAwKQogICAgICByZXMuYSA9IGFiLmI7CiAgICBlbHNlCiAgICB7CiAgICAgIHN0ZDo6bnRoX2VsZW1lbnQocmFua3MubC5iZWdpbigpICsgMSwKCQkgICAgICAgcmFua3MubC5lbmQoKSAtIHIxLAoJCSAgICAgICByYW5rcy5sLmVuZCgpKTsKICAgICAgcmVzLmEgPSAqKHJhbmtzLmwuZW5kKCkgLSByMSk7CiAgICB9CiAKICAgIGlmIChyYW5rcy5yYV9sZXNzIDw9IGsyIC0gMSkKICAgICAgcmVzLmIgPSBhYi5hOwogICAgZWxzZSBpZiAocjIgPD0gMCkKICAgICAgcmVzLmIgPSBhYi5iOwogICAgZWxzZQogICAgewogICAgICBzdGQ6Om50aF9lbGVtZW50KHJhbmtzLmwuYmVnaW4oKSArIDEsCgkJICAgICAgIHJhbmtzLmwuZW5kKCkgLSByMiwKCQkgICAgICAgcmFua3MubC5lbmQoKSk7CiAgICAgIHJlcy5iID0gKihyYW5rcy5sLmVuZCgpIC0gcjIpOwogICAgfQogfQoKICByZXR1cm4gcmVzOwp9CgppbnQgc2VsZWN0KGNvbnN0IFZlYyYgYTEsIGNvbnN0IFZlYyYgYTIsIGludCBrKQp7CiAgYXNzZXJ0KGExLnNpemUoKSA9PSBhMi5zaXplKCkpOwogIEFCIGFiID0gYmlzZWxlY3QoYTEsIGEyLCBrLCBrKTsKICBhc3NlcnQoYWIuYSA9PSBhYi5iKTsKICByZXR1cm4gYWIuYTsKfQoKaW50IGNoZWNrKGNvbnN0IFZlYyYgYTEsIGNvbnN0IFZlYyYgYTIsIGludCBrKQp7CiAgaW50IG4gPSBhMS5zaXplKCkgLSAxOwogIFZlYyBzdW07CiAgc3VtLnJlc2VydmUobiAqIG4pOwoKICBmb3IgKGludCBpID0gMTsgaSA8PSBuOyArK2kpCiAgICBmb3IgKGludCBqID0gMTsgaiA8PSBuOyArK2opCiAgICAgIHN1bS5wdXNoX2JhY2soYTFbaV0gKyBhMltqXSk7CgogIHN0ZDo6bnRoX2VsZW1lbnQoc3VtLmJlZ2luKCksCgkJICAgc3VtLmVuZCgpIC0gaywKCQkgICBzdW0uZW5kKCkpOwogIHJldHVybiAqKHN1bS5lbmQoKSAtIGspOwp9CgojaW5jbHVkZSA8cmFuZG9tPgoKaW50IG1haW4oKQp7CiAgc3RkOjpyYW5kb21fZGV2aWNlIHJkOwogIHN0ZDo6ZGVmYXVsdF9yYW5kb21fZW5naW5lIGUxKHJkKCkpOwogIHN0ZDo6dW5pZm9ybV9pbnRfZGlzdHJpYnV0aW9uPGludD4gdW5pZm9ybV9kaXN0KDEsIDYpOwoKICBpbnQgbiA9IDEwMDAwMDsKICBWZWMgYTEsIGEyOwogIGExLnJlc2VydmUobisxKTsKICBhMi5yZXNlcnZlKG4rMSk7CiAgaW50IHggPSAxMDAwOwoKICBhMS5wdXNoX2JhY2soMCk7CiAgZm9yIChpbnQgaSA9IDA7IGkgPCBuOyArK2kpCiAgewogICAgeCAtPSB1bmlmb3JtX2Rpc3QoZTEpOwogICAgYTEucHVzaF9iYWNrKHgpOwogIH0KCiAgeCA9IDEwMDA7CgogIGEyLnB1c2hfYmFjaygwKTsKICBmb3IgKGludCBpID0gMDsgaSA8IG47ICsraSkKICB7CiAgICB4IC09IHVuaWZvcm1fZGlzdChlMSk7CiAgICBhMi5wdXNoX2JhY2soeCk7CiAgfQoKICBzdGQ6OmNvdXQgPDwgc2VsZWN0KGExLCBhMiwgMSArIG4gKiBuIC8gNykgPDwgJ1xuJwogICAgLyo8PCBjaGVjayhhMSwgYTIsIDEgKyBuICogbiAvIDcpIDw8ICdcbicqLzsKCiAgcmV0dXJuIDA7Cn0K