Finding the kth-smallest element in union of sorted arrays

algorithm, python

Solution

We make an observation that when `Ai < Bj`, then it must be true that `Ai < Bj-1`. On the other hand, if `Bj < Ai`, then `Bj < Ai-1`.. How can it be true for any `i` and `j`?

It isn't true for all pairs of `i` and `j`. The article considers a special situation.

First, it is assumed that there are no duplicates, not even in the form of common elements of `A` and `B`. Second, the conclusion that

Ai < Bj ==> Ai < Bj-1,   resp.  Bj < Ai ==> Bj < Ai-1

is made under the condition that neither of

Bj-1 < Ai < Bj  resp. Ai-1 < Bj < Ai

holds. So by excluding these configurations, `Ai < Bj ==> Ai <= Bj-1` and `Bj < Ai ==> Bj <= Ai-1` follow immediately, and the strict inequalities then follow by the assumption that no duplicates exist.

We try to approach this tricky problem by comparing middle elements of A and B, which we identify as Ai and Bj. If Ai is between Bj and Bj-1, we have just found the i+j+1 smallest element

In array `B`, there are `j` elements smaller than `Bj`, and in array `A`, there are `i` elements smaller than `Ai` (indices start at 0). So if `Bj-1 < Ai < Bj`, both arrays together contain exactly `j + i` elements that are smaller than `Ai`.

What changes if there are duplicates?

Not much.

We still consider the situation where `i + j = k-1`. Let us assume that `Ai <= Bj`.

- What if `Ai = Bj`?

- What if `Ai < Bj`?

In case 1., let `m` be the smallest index such that `Am = Ai`, and `n` the smallest index such that `Bn = Bj`. Then in both arrays together, there are exactly `m + n <= i + j = k-1` elements strictly smaller than `Ai`, and at least `(i+1) + (j+1) = (k+1)` elements not larger than `Ai`. Hence the k-th smallest element is equal to `Ai`.

For 2., we have three cases to consider, a) `Bj-1 < Ai`, b) `Bj-1 = Ai`, c) `Bj-1 > Ai`.

In case a), we have `j` elements in `B` that are not larger than `Ai`, and they are all strictly smaller, and we have `m <= i` elements in `A` that are strictly smaller than `Ai` (`m` as above) and an unkown number, but at least `i-m+1` elements equal to `Ai`. So there are exactly `j + m <= j + i = k-1` elements in both arrays together that are strictly smaller than `Ai`, and at least `j + m + (i-m+1) = j+i+1 = k` elements not larger than `Ai`, hence the k-th smallest element of both arrays together is equal to `Ai`.

In case b), the same reasoning shows that the k-th smallest element of both arrays together is equal to `Ai`.

In the remaining case, `Ai < Bj-1`, things become hardly more complicated. Array `B` contains at least `j` elements not larger than `Bj-1`, and array `A` contains at least `i+1` elements strictly smaller than `Bj-1`, hence the k-th smallest element of both arrays together is at most as large as `Bj-1`. But it cannot be smaller than `Ai` (`B` contains at most `j-1` elements smaller than `Ai`, so both arrays together contain at most `i + (j-1) = k-2` elements smaller than `Ai`).

So we can still discard the part below `Ai` from the array `A` and the part above `Bj-1` from the array `B` and proceed as without duplicates.

All that changed was that a few strict inequalities had to be replaced with weak inequalities.

The code (would be more efficient if starting indices and lengths were passed instead of slicing, but slicing yields shorter code):

def kthsmallest(A, B, k):
    if k < 1:
        return None
    a_len, b_len = len(A), len(B)
    if a_len == 0:
        return B[k-1] # let it die if B is too short, I don't care
    if b_len == 0:
        return A[k-1] # see above
    # Handle edge case: if k == a_len + b_len, we would
    # get an out-of-bounds index, since i + j <= a_len+b_len - 2
    # for valid indices i and j
    if a_len + b_len == k:
        if A[-1] < B[-1]:
            return B[-1]
        else:
            return A[-1]
    # Find indices i and j approximately proportional to len(A)/len(B)
    i = (a_len*(k-1)) // (a_len+b_len)
    j = k-1-i
    # Make sure the indices are valid, in unfortunate cases,
    # j could be set to b_len by the above
    if j >= b_len:
        j = b_len-1
        i = k-1-j
    if A[i] <= B[j]:
        if j == 0 or B[j-1] <= A[i]:
            return A[i]
        # A[i] < B[j-1] <= B[j]
        return kthsmallest(A[i:], B[:j], k-i)
    # B[j] < A[i], symmetrical to A[i] < B[j]
    if i == 0 or A[i-1] <= B[j]:
        return B[j]
    # B[j] < A[i-1]
    return kthsmallest(A[:i], B[j:], k-j)

Problem

Original source