Finding the kth-smallest element in union of sorted arrays

Question

I was studying the article on finding the kth-smallest element in the union of two sorted arrays at leetcode . I don't think that the algorithm is correct. There is this line: 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 ?

And secondly, this line also baffles me: 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 , although it comes out to be true. Can anyone explain the reason? I really want to understand the algo, I have done it by merging the arrays, but that takes O(N) time, compared to O(log N) time here.

Answer 1

You're interpreting these statements in isolation, but they build on one another. Here's the text that (I think) you're referring to:

Maintaining the invariant i + j = k – 1, If Bj-1 < Ai < Bj, then Ai must be the k-th smallest, or else if Ai-1 < Bj < Ai, then Bj must be the k-th smallest. If one of the above conditions are satisfied, we are done. If not, we will use i and j as the pivot index to subdivide the arrays. But how? Which portion should we discard? How about Ai and Bj itself?

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. Why?

Breaking this down into sub-propositions yields the following interpretation (keeping in mind that indexing starts from 0 , but that A0 is the first smallest item, and A1 is the second smallest item, and so on):

1. i + j = k - 1 (invariant, by definition)
2. Posit that Bj-1 < Ai < Bj . Then Ai must be the k th smallest. This is because Ai is greater than i items in A and is greater than j items in B . So it's greater than a total of i + j = k - 1 items. That means its index in a merged A|B list would be k - 1 , and so it would be the k th item in that list.
3. Posit that Ai-1 < Bj < Ai . Then Bj must be the k th smallest, by the same line of reasoning in 2.
4. Now posit that both (a) Bj-1 < Ai < Bj and (b) Ai-1 < Bj < Ai are false . Then it follows quite obviously that if Ai < Bj then A1 < Bj-1 , because otherwise (a) would be true. And likewise, if Bj < Ai then Bj < Ai-1 , because otherwise, (b) would be true.

I'm taking you at your word that you want an explanation of these statements rather than of the algorithm as a whole. (But I'll say more if you like.)

Note also that, as Daniel Fischer 's answer reminds me, the above reasoning only holds if there are no duplicates; call that proposition 0.

Answer 2

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 .

1. What if Ai = Bj ?
2. 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)