[LeetCode] 373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.
Define a pair (u,v) which consists of one element from the first array and one element from the second array.
Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.
Example 1:
Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
Example 2:
Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
Example 3:
Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

Thought process:
Maintain a min heap of size k. At first, offer all (nums1[i], nums2[0]) to the heap. Do k iterations. For each iteration, poll one pair from the heap, add it to the result list, and offer (nums1[i], nums2[j + 1]) to the heap.

Solution:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> pairs = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
            return pairs;
        }
        
        Queue<int[]> heap = new PriorityQueue<>((a, b) -> a[0] + a[1] - b[0] - b[1]);
        for (int i = 0; i < k && i < nums1.length; i++) {
            heap.offer(new int[]{ nums1[i], nums2[0], 0 });
        }
        
        while (k > 0 && !heap.isEmpty()) {
            int[] smallest = heap.poll();
            pairs.add(new int[]{ smallest[0], smallest[1] });
            
            int index = smallest[2];
            if (index < nums2.length - 1) {
                heap.offer(new int[]{ smallest[0], nums2[index + 1], index + 1 });
            }
            k--;
        }
        return pairs;
    }
}

Time complexity:
There're at most k iterations for both the for loop and the while loop. The overall time complexity is O(klogk).

Comments

Popular posts from this blog

[LeetCode] 631. Design Excel Sum Formula

[LeetCode] 714. Best Time to Buy and Sell Stock with Transaction Fee

[LeetCode] 269. Alien Dictionary