The k^th Greatest Element In Two Sorted Arrays

Given two sorted arrays A and B, find the k^th greatest element of all the elements.

Take the i^th element from array A and the j^th element from array B so that i + j == k, we can see that if B[j] <= A[i] <= B[j+1], then A[i] must be the k^th element of all elements in A and B. So now we have 3 scenarios to consider:

1. A[i] == B[j], then either A[i] or B[j] is the k^th element
2. A[i] < B[j], then the k^th element must be in A[i+1..A.length] or B[1..j]
3. A[i] > B[j], this is symmetric to case 2

Case 1 is obvious, for case 2, if A[i-p] is the k^th element, then A[i-p] must be greater than or equal to j + p elements in B, so A[i-p] >= B[j+p], since A[i] < B[j], then A[i-p] < B[j] < B[j+p], which contradicts to A[i-p] >= B[j+p], so the k^th element must be in A[i+1..A.length]. Same proof can be used to show that the k^th element must be in B[1..j].

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// assume that k is always less than A.length + B.length
def k_th_element(A, B, k)
    if A.length > B.length:
        swap(A, B) // make sure A is always the smaller one
    if A.length == 0:
        return B[k]
    if k == 1:
        return min(A[0], B[0])
    i = min(A.length, k / 2)
    j = k - i
    if A[i] == B[j]:
        return A[i]
    else if A[i] < B[j]:
        return k_th_element(A[i+1..A.length], B[1..j], k - i)
    else:
        return k_th_element(A[1..i], B[j+1..B.length], k - j)