Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Find k-th element in two sorted arrays
// Problem:
// For two sorted arrays A and B and index k find an element on k-th position for merged
// A+B array (sorting preserved).
// Limitations:
// - Complexiti is lower than O(n)
// - No memory allocation
// - Input data can't be modified.
//
// Proposed solution:
// For index i in A and index j in B such that i + j == k move indexes by half of their
// values (one goes left and another right) until B[j-1]<A[i]<B[j] or A[i-1]<B[j]<A[i]
// found. The element in the middle is the k-th element.
#include <algorithm>
#include <iostream>
#include <iterator>
#include <stdexcept>
#include <vector>
using namespace std;
// set to true for iterations count output
bool countIterations = false;
size_t iterations = 0;
int findKthInTwoSortedArrays(const int *a, size_t m, const int *b, size_t n, size_t k)
{
if (k >= m + n) throw invalid_argument("k");
iterations = 0;
if (0 == m) return b[k];
if (0 == n) return a[k];
if (0 == k) return min(a[0], b[0]);
if (m + n == k) return max(a[m - 1], b[n - 1]);
size_t i = 0, j = 0, delta = 0;
if (m < n)
{
i = min(k / 2, m - 1);
delta = i / 2;
j = k - i;
}
else
{
j = min(k / 2, n - 1);
delta = j / 2;
i = k - j;
}
delta += delta ? 0 : 1;
// always keep indexes values to satisfy i + j == k - 1
while (1)
{
++iterations;
bool iIsInRange = i < m,
iHasItemsOnLeft = i > 0,
jIsInRange = j < n,
jHasItemsOnLeft = j > 0;
// when a[i] sits exactly between b[j-1] and b[j]
// it is the kth element
if (iIsInRange)
{
if ((!jHasItemsOnLeft || (b[j - 1] <= a[i])) &&
(!jIsInRange || (a[i] <= b[j])))
{
return a[i];
}
}
// when b[j] sits exactly between a[i-1] and a[i]
// it is the kth element
if (jIsInRange)
{
if ((!iHasItemsOnLeft || (a[i - 1] <= b[j])) &&
(!iIsInRange || (b[j] <= a[i])))
{
return b[j];
}
}
if (jIsInRange && (!iIsInRange || (b[j] <= a[i])))
{
// b[j] is smaller both a[i] and a[i-1]
if (j == n - 1)
{
// reached rightmost element of b[]
return a[k - j - 1];
}
// move i left and j right
if (j + delta >= n) delta = n - j - 1;
if (delta > i) delta = i;
j += delta;
i -= delta;
}
else
{
if (i == m - 1)
{
return b[k - i - 1];
}
if (i + delta >= m) delta = m - i - 1;
if (delta > j) delta = j;
i += delta;
j -= delta;
}
delta /= 2;
delta += delta ? 0 : 1;
}
}
struct TestData
{
vector<int> a, b;
};
bool testResult(const TestData& testData)
{
const auto&a = testData.a;
const auto&b = testData.b;
size_t mergedSize = a.size() + b.size();
vector<int> merged(mergedSize);
merge(a.cbegin(), a.cend(), b.cbegin(), b.cend(), merged.begin());
for (size_t i = 0; i < mergedSize; ++i)
{
cout << "\tindex " << i + 1;
auto expected = merged[i];
auto aSize = a.size();
auto bSize = b.size();
auto actual =
findKthInTwoSortedArrays(aSize ? &a[0] : nullptr, aSize,
bSize ? &b[0] : nullptr, bSize,
i);
if (actual != expected)
{
cout << " failed!" << endl;
cerr << "Expected=" << expected << ", actual=" << actual << ", index=" << i <<
"\nArrays:\n";
copy(a.cbegin(), a.cend(), ostream_iterator<int>(cerr, ", "));
cerr << endl;
copy(b.cbegin(), b.cend(), ostream_iterator<int>(cerr, ", "));
cerr << endl;
return false;
}
cout << " succeeded!" << endl;
if (countIterations)
{
cout << "n=" << mergedSize << ", iterations=" << iterations << endl;
}
}
cout << endl;
return true;
}
int main()
{
TestData tests[] =
{
{ {}, { 1, 2, 3 } },
{ { 1, 2 }, {} },
{ { 2 }, { 42 } },
{ { 2 }, { 1, 3 } },
{ { 1, 3 }, { 2 } },
{ { 1, 2, 3 }, { 10, 20, 30, 40 } },
{ { 10, 20, 30 }, { 1, 2 } },
{ { -1, 0, 1 }, { -2, 2 } },
{ { 2, 4, 8, 16 }, { 1, 3, 5, 9, 11, 13, 17 } },
{ { 1, 10 }, { 2, 3, 4, 5, 6, 7, 8, 9 } },
{ { 1, 3, 4, 5, 6, 7, 9 }, { 2, 8 } },
{
{ 1, 2, 3, 55, 56, 78, 79, 79, 79, 100, 101, 1000, 5000, 9001},
{ 4, 5, 6, 7, 8, 9, 11, 13, 17, 19, 23, 27, 29, 37, 99, 102, 2000, 2001 }
},
{ { 10 }, { 1, 2, 3, 4, 12, 13, 14, 15 } },
{ { 67, 122, 135, 156 }, { 11, 15, 26, 66, 111, 255 } },
{
{ 0, 0, 1, 3, 4, 4, 7, 7, 10, 10, 11, 11, 11, 12, 13, 16, 16 },
{ 2, 2, 2, 4, 5, 5, 6, 7, 8, 9, 9, 10, 10, 10, 11, 13, 13, 15, 16, 16, 17, 17, 18, 18, 18, 18, 20, 20, 20, 21, 21, 22, 24, 26, 26, 27, 29, 32, 33, 35, 36, 42, 42 }
},
{ { 1, 3, 4, 5, 5, 6, 8, 9, 9, 10, 11, 11 }, { 1, 1 } },
{
{ 0, 1, 1, 1, 1, 3, 3, 7 },
{ 0, 0, 1, 1, 1, 4, 4, 7, 9, 10, 10, 12, 13, 16, 16, 16, 17, 19, 19, 20, 20, 20, 23, 23, 25, 25, 25, 26, 28, 28, 28, 29, 29, 30, 33, 33, 35, 35, 35, 37, 38 }
},
{ { 1, 1 }, { 2, 2, 3, 4, 6, 6, 7, 7, 8 } },
{ { 0 }, { 0, 0, 1, 2, 2, 5 } },
{ { 0, 1, 3, 3, 4, 5, 6 }, { 1, 1, 2, 2, 3, 3, 4, 6 } },
{ { 18, 19, 19, 19, 19, 20 }, { 15, 16, 17, 19, 19, 19 } },
{
{ 0, 0, 1, 1, 1, 1, 3, 3, 5, 10, 10, 11, 16, 16, 17, 18, 19, 19, 19, 19, 20 },
{ 3, 3, 3, 5, 5, 5, 5, 7, 7, 8, 8, 10, 11, 13, 13, 15, 16, 17, 19, 19, 19 }
},
{ { 0, 2, 2, 2, 3, 4 }, { 0, 1, 2, 2, 3, 4, 5, 6, 6, 8 } },
{
{ 0, 0, 2, 2, 3, 3, 4, 6, 7, 8, 9, 10, 12, 12, 13, 15, 15, 16, 16, 18, 18, 18, 21, 21, 24, 25 },
{ 0, 2, 2, 5, 6, 6, 6, 7, 7, 8, 8, 8, 8, 13, 14, 14, 14, 14, 19, 20, 20 }
}
};
size_t succeeded = 0, failed = 0;
int i = 1;
for (const auto& t : tests)
{
cout << "Test set " << i++ << endl;
if (!testResult(t))
{
++failed;
cerr << "\nFix it!!11" << endl;
break;
}
else
{
++succeeded;
}
}
cout << "Total " << failed + succeeded << " tests: " << failed << " failed, " <<
succeeded << " succeeded." << endl;
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment