Kth Smallest element

Q: Given two sorted arrays A, B of size m and n respectively. Find the k-th smallest element in the union of A and B. You can assume that there are no duplicate elements.

One approach we could think of is merge both arrays and the k-th smallest element could be accessed directly. Merging would require extra space of O(m+n). The linear run time is pretty good, but could we improve it even further?

Another approach is using two pointers, you can traverse both arrays without actually merging them, thus without the extra space. Both pointers are initialized to point to head of A and B respectively, and the pointer that has the larger of the two is incremented one step. The k-th smallest is obtained by traversing a total of k steps.

We can solve this problem in logarithmic time in O(lg m + lg n) using binary search approach.

We try to approach this tricky problem by comparing middle elements of A and B, which we identify as A_i and B_j. If A_i is between B_j and B_j-1, we have just found the i+j+1 smallest element. Therefore, if we choose i and j such that i+j = k-1, we are able to find the k-th smallest element. This is an important invariant that we must maintain for the correctness of this algorithm.

Maintaining the invariant
    i + j = k – 1,

If B_j-1 < A_i < B_j, then A_i must be the k-th smallest,
or else if A_i-1 < B_j < A_i, then B_j must be the k-th smallest.

We can make an observation that when A_i < B_j, then it must be true that A_i < B_j-1. On the other hand, if B_j < A_i, then B_j < A_i-1.

Using the above relationship, it becomes clear that when A_i < B_j, A_i and its lower portion could never be the k-th smallest element. So do B_j and its upper portion. Therefore, we could conveniently discard A_i with its lower portion and B_j with its upper portion.