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

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)