[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
Post a Comment