The k-th smallest of two sorted arrays
Given two arrays a and b, each sorted in non-decreasing order (either may be
empty, but not both), and an integer k with 1 ≤ k ≤ |a| + |b|, return the k-th smallest
value of the two arrays taken together. Use O(log k) steps, without merging.
Examples
| Input | Output | Why |
|---|---|---|
a = [2, 3, 6, 7, 9], b = [1, 4, 8, 10], k = 5 | 6 | Merged: 1, 2, 3, 4, 6, 7, 8, 9, 10. |
a = [1, 3], b = [2], k = 2 | 2 | |
a = [], b = [5, 9], k = 1 | 5 |
Constraints
- 0 ≤ |a|, |b| ≤ 10⁵, |a| + |b| ≥ 1
- 1 ≤ k ≤ |a| + |b|
- −10⁹ ≤ values ≤ 10⁹
Hints
Hint 1
Merging until the k-th element costs O(k). The goal is to throw away many candidates at a time.
Hint 2
Take h = k / 2 elements from the front of each array and compare the last of each: a[h − 1] and b[h − 1].
Hint 3
If a[h − 1] ≤ b[h − 1], at most 2h − 2 elements can come before a[h − 1], so it ranks below k, and so does everything before it in a. Drop those h elements and look for the (k − h)-th of what is left.
Where people lose marks
- Taking h elements from an array that has fewer left: clamp the block to what remains.
- Discarding the block of the larger element: it may hold the answer.
- Forgetting k = 1, where h = 0 and nothing would be discarded: the loop never ends.
Dry run on [2,3,6,7,9,12,15],[1,4,8,10,11,14],8
Read it one frame at a time, the way you would trace the code on paper in an interview. Each frame shows the state after the step it describes.
-
Step 1 / 5k8
Top row a, bottom row b. We want the 8-th smallest of the two together.
-
Step 2 / 5k8h4a block end7b block end10
k = 8, so take up to 4 from each. 7 ≤ 10: the 4 elements of a up to 7 all rank below 8. Drop them; k becomes 4.
-
Step 3 / 5k4h2a block end12b block end4
k = 4, so take up to 2 from each. 12 > 4: the 2 elements of b up to 4 all rank below 4. Drop them; k becomes 2.
-
Step 4 / 5k2h1a block end9b block end8
k = 2, so take up to 1 from each. 9 > 8: the 1 element of b up to 8 all rank below 2. Drop them; k becomes 1.
-
Step 5 / 5k1
k = 1: the answer is the smaller front element, min(9, 10) = 9.
How to think about it
Keep a start index into each array, i and j, and the rank k still
wanted among the elements from there on. Take up to h = ⌊k/2⌋ elements from the front of each, and
compare the last element of each block, x = a[i + h − 1] and y = b[j + h − 1].
Suppose x ≤ y. How many remaining elements can be smaller than x? At most the
h − 1 before it in a, and at most the h − 1 before y in
b. So x ranks at most 2h − 1 ≤ k − 1: it and everything before it in
a rank below k, and the h elements a[i..i + h − 1] can be discarded:
i += h, k −= h. The case y < x is symmetric.
Three stopping cases: one array is used up (the answer is in the other, directly indexed), or
k = 1 (the answer is the smaller front element). If an array has fewer than h elements
left, take all of them as its block; the argument is unchanged.
k at least halves with every step, so there are O(log k) steps. With k = (m + n + 1)/2 this gives the median of two sorted arrays; the next problem does that with a different and faster-to-code search over where to cut the shorter array.
Complexity. Time O(log k) · Space O(1)
The solution, in six languages
Each one is compiled and run against the test table before it is published. Java is reviewed by hand.
def kth_of_two(a, b, k):
i = j = 0 # a[i:] and b[j:] are still in play; want their k-th
while True:
if i == len(a):
return b[j + k - 1]
if j == len(b):
return a[i + k - 1]
if k == 1:
return min(a[i], b[j])
h = k // 2
ni = min(i + h, len(a)) - 1 # last index of a's block
nj = min(j + h, len(b)) - 1
if a[ni] <= b[nj]:
k -= ni - i + 1 # a[i..ni] all rank below k
i = ni + 1
else:
k -= nj - j + 1
j = nj + 1#include <vector>
#include <algorithm>
int solve(const std::vector<int>& a, const std::vector<int>& b, int k) {
const int m = (int)a.size(), n = (int)b.size();
int i = 0, j = 0; // a[i..] and b[j..] are in play
for (;;) {
if (i == m) return b[j + k - 1];
if (j == n) return a[i + k - 1];
if (k == 1) return std::min(a[i], b[j]);
int h = k / 2;
int ni = std::min(i + h, m) - 1, nj = std::min(j + h, n) - 1;
if (a[ni] <= b[nj]) { k -= ni - i + 1; i = ni + 1; } // drop a's block
else { k -= nj - j + 1; j = nj + 1; } // drop b's block
}
}class Solution {
public int kthOfTwo(int[] a, int[] b, int k) {
int m = a.length, n = b.length, i = 0, j = 0; // a[i..] and b[j..] are in play
while (true) {
if (i == m) return b[j + k - 1];
if (j == n) return a[i + k - 1];
if (k == 1) return Math.min(a[i], b[j]);
int h = k / 2;
int ni = Math.min(i + h, m) - 1, nj = Math.min(j + h, n) - 1;
if (a[ni] <= b[nj]) { k -= ni - i + 1; i = ni + 1; } // drop a's block
else { k -= nj - j + 1; j = nj + 1; } // drop b's block
}
}
}function kthOfTwo(a, b, k) {
let i = 0, j = 0; // a[i..] and b[j..] are in play
for (;;) {
if (i === a.length) return b[j + k - 1];
if (j === b.length) return a[i + k - 1];
if (k === 1) return Math.min(a[i], b[j]);
const h = k >> 1;
const ni = Math.min(i + h, a.length) - 1, nj = Math.min(j + h, b.length) - 1;
if (a[ni] <= b[nj]) { k -= ni - i + 1; i = ni + 1; } // drop a's block
else { k -= nj - j + 1; j = nj + 1; } // drop b's block
}
}fn solve(a: &[i64], b: &[i64], k: i64) -> i64 {
let (mut i, mut j, mut k) = (0usize, 0usize, k as usize); // a[i..] and b[j..] in play
loop {
if i == a.len() {
return b[j + k - 1];
}
if j == b.len() {
return a[i + k - 1];
}
if k == 1 {
return a[i].min(b[j]);
}
let h = k / 2;
let ni = (i + h).min(a.len()) - 1;
let nj = (j + h).min(b.len()) - 1;
if a[ni] <= b[nj] {
k -= ni - i + 1; // drop a's block
i = ni + 1;
} else {
k -= nj - j + 1; // drop b's block
j = nj + 1;
}
}
}import Data.Array
solve :: [Int] -> [Int] -> Int -> Int
solve xs ys = go 0 0
where
(m, n) = (length xs, length ys)
a = listArray (0, m - 1) xs
b = listArray (0, n - 1) ys
-- a[i..] and b[j..] are in play; want their k-th smallest
go i j k
| i == m = b ! (j + k - 1)
| j == n = a ! (i + k - 1)
| k == 1 = min (a ! i) (b ! j)
| a ! ni <= b ! nj = go (ni + 1) j (k - (ni - i + 1)) -- drop a's block
| otherwise = go i (nj + 1) (k - (nj - j + 1)) -- drop b's block
where
h = k `div` 2
ni = min (i + h) m - 1
nj = min (j + h) n - 1Variants to try
- The median of two sorted arrays: the average of the k-th and (k + 1)-th for even total length.
- The k-th smallest of many sorted arrays: a heap of their fronts, O(k log r), or a binary search on the value.
Next problems
Median of two sorted arrays
Choose how many elements the left half takes from the shorter array; the rest come from the longer one. One comparison says whether to take more or fewer.
Open the problem → Easy10 minFind a value in a sorted array
The search every other binary search is built on: halve a closed range until the value is found or the range is empty.
Open the problem → Easy10 minInteger square root without a library
Search the candidates 0 to n for the last x with x² ≤ n, and keep the square from overflowing.
Open the problem →Write a solution and run it against the real test table.
Keep going
Median of two sorted arrays
Choose how many elements the left half takes from the shorter array; the rest come from the longer one. One comparison says whether to take more or fewer.
Open the problem → Easy10 minFind a value in a sorted array
The search every other binary search is built on: halve a closed range until the value is found or the range is empty.
Open the problem → Easy10 minInteger square root without a library
Search the candidates 0 to n for the last x with x² ≤ n, and keep the square from overflowing.
Open the problem → Binary searchThe pattern behind it
Search the answer, not the array: halve a monotone predicate until one candidate is left.
Read the pattern →Preparing for a real process? Quant interview preparation, or book a free 20-minute call.