Find k-th element in two sorted arrays

findKthElementInTwoSortedArrays.cpp

// 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;
}