Find kth smallest pair sum with heaps
Company: Meta
Role: Machine Learning Engineer
Category: Coding & Algorithms
Difficulty: Medium
Interview Round: Onsite
Quick Answer: This question evaluates algorithmic problem-solving and complexity-analysis skills, focusing on efficient selection from two sorted arrays while handling duplicate sums and boundary conditions; it belongs to the Coding & Algorithms domain.
Constraints
- 1 <= n, m
- 1 <= k <= n * m
- A and B are sorted in nondecreasing order
- All elements of A and B are non-negative integers
- Do not generate all n*m pair sums explicitly
Examples
Input: ([1, 7, 11], [2, 4, 6], 3)
Expected Output: 7
Explanation: All pair sums sorted: [3, 5, 7, 9, 11, 13, 13, 15, 17]. The 3rd smallest is 7 (= 1 + 6).
Input: ([1, 1, 2], [1, 2, 3], 1)
Expected Output: 2
Explanation: k=1 boundary: the smallest sum is A[0] + B[0] = 1 + 1 = 2.
Input: ([1, 2], [3], 2)
Expected Output: 5
Explanation: Sums are [4, 5]; the 2nd smallest is 5 (= 2 + 3).
Input: ([0, 0], [0, 0], 4)
Expected Output: 0
Explanation: All four pair sums equal 0; duplicates count with multiplicity, so the 4th smallest is still 0.
Input: ([1, 7, 11], [2, 4, 6], 9)
Expected Output: 17
Explanation: k = n*m = 9 boundary: the largest pair sum 11 + 6 = 17.
Input: ([5], [10], 1)
Expected Output: 15
Explanation: Single element each: the only sum is 5 + 10 = 15.
Input: ([2, 4, 6], [1, 3, 5], 6)
Expected Output: 7
Explanation: Sums sorted: [3, 5, 5, 7, 7, 7, 9, 9, 11]. The 6th smallest is 7, showing ties are counted individually.
Hints
- The smallest pair sum is always A[0] + B[0]. The next smallest is one step right in either array — this 'frontier' is exactly what a min-heap tracks efficiently.
- Seed the heap with one entry per row of the SHORTER array: (smaller[i] + larger[0], i, 0). This caps the heap at O(min(n, m)) entries.
- When you pop (sum, i, j), the only new candidate it unlocks is (smaller[i] + larger[j+1], i, j+1). Push it if j+1 is in bounds. Never push a 'down' move, or you would enqueue the same cell twice.
- Duplicate sums count with multiplicity: popping k times naturally handles ties because each (i, j) cell is a distinct heap entry even when sums are equal.