K-th Element of Two Sorted Arrays

The problem, simply

You get two arrays, each sorted from small to big, of sizes n and m, and a number k. Imagine pouring both into one sorted array. Return the number at position k (counting from 1).

You should not build the big array. The fast answer looks at only a few numbers.

A story to remember it

Picture it

Examples

Way 1: the simple way

Do the merge step of merge sort, but stop as soon as you have taken k numbers. Two fingers start at the front of each array. Each step, take the smaller front number and move that finger.

class Solution {
    public int kthElement(int[] a, int[] b, int k) {
        int i = 0, j = 0, last = -1;
        for (int taken = 0; taken < k; taken++) {
            if (j == b.length || (i < a.length && a[i] <= b[j])) {
                last = a[i++];          // a's front is smaller (or b is used up)
            } else {
                last = b[j++];
            }
        }
        return last;                    // the k-th number taken
    }
}

How long it takes: k steps, O(k), which can be n + m. Extra room: O(1).

Way 2: the clever way

The k smallest numbers are always a front part of a plus a front part of b. Say we take c1 from a and c2 = k - c1 from b. Binary search for the right c1, on the shorter array.

1. Make a the shorter array. c1 can be from max(0, k - m) to min(k, n).
2. Guess c1 in the middle; c2 = k - c1.
3. Look at the border: l1 = a[c1 - 1] and l2 = b[c2 - 1] are the last ones taken; r1 = a[c1] and r2 = b[c2] are the first ones left. Use minus and plus infinity when a side is empty.
4. If l1 <= r2 and l2 <= r1, the split is right. The answer is max(l1, l2).
5. If l1 > r2, we took too many from a: hi = c1 - 1.
6. Otherwise we took too few from a: lo = c1 + 1.

The bottom row in the animation lists every possible c1.

class Solution {
    public int kthElement(int[] a, int[] b, int k) {
        if (a.length > b.length) return kthElement(b, a, k);   // search the shorter one
        int n = a.length, m = b.length;
        int lo = Math.max(0, k - m), hi = Math.min(k, n);
        while (lo <= hi) {
            int c1 = lo + (hi - lo) / 2;         // how many from a
            int c2 = k - c1;                     // the rest from b
            int l1 = c1 == 0 ? Integer.MIN_VALUE : a[c1 - 1];
            int l2 = c2 == 0 ? Integer.MIN_VALUE : b[c2 - 1];
            int r1 = c1 == n ? Integer.MAX_VALUE : a[c1];
            int r2 = c2 == m ? Integer.MAX_VALUE : b[c2];
            if (l1 <= r2 && l2 <= r1) return Math.max(l1, l2);   // a clean border
            if (l1 > r2) hi = c1 - 1;            // too many from a
            else lo = c1 + 1;                    // too few from a
        }
        return -1;                               // only if k is out of range
    }
}

How long it takes: the range of c1 has at most min(n, m) + 1 choices and halves each time, O(log(min(n, m))). Extra room: O(1).

Watch out for

• Forgetting the empty sides. When c1 == 0, nothing was taken from a, so l1 must not block anything: use the smallest possible value. When c1 == n, nothing is left in a: use the biggest possible value.
• Searching on the longer array. It still works if the range is right, but it is slower and the range math is easier to get wrong.
• k counts from 1 here. If your judge counts from 0, add 1 first.

What interviewers ask next

• Median of two sorted arrays (LeetCode 4)? It is this problem with k = (n + m + 1) / 2, plus the next number when n + m is even, which is min(r1, r2) at the same split.
• Another O(log k) way? Compare a[k/2 - 1] and b[k/2 - 1]. The smaller side's first k/2 numbers cannot hold the answer, so drop them and look for the (k - k/2)-th in what is left.
• What about the k-th smallest in a sorted matrix, or among many sorted lists? Use a min heap of fronts (O(k log lists)), or binary search on the value and count how many numbers are at most it.
• Why is max(l1, l2) the answer? The taken part holds exactly k numbers, all at most everything left. So the biggest taken number is in position k.

Remember this

Grown-up words: the guess is a partition (or cut) of each array, the border check is the partition condition, minus and plus infinity are sentinels, and the median is the special case k = (n + m + 1) / 2.