[刷题防痴呆] 0373 - 查找和最小的K对 (Find K

题目地址

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/

题目描述

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:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: 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:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: 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:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [1,3],[2,3]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

思路

• minHeap.
• 本质上是merge sort. 用heap实现.
• heap中维护一个数组, nums1, nums2, 在nums2中的index.

关键点

代码

• 语言支持：Java

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(
        (a, b) -> (a[0] + a[1]) - (b[0] + b[1]));
        List<List<Integer>> ans = new ArrayList<>();
        
        int length1 = nums1.length;
        int length2 = nums2.length;
        
        if (length1 == 0 || length2 == 0) {
            return ans;
        }
        for (int i = 0; i < Math.min(length1, k); i++) {
            pq.offer(new int[]{nums1[i], nums2[0], 0});
        }

        while (!pq.isEmpty() && k > 0) {
            int[] cur = pq.poll();
            List<Integer> curAns = new ArrayList<>();
            curAns.add(cur[0]);
            curAns.add(cur[1]);
            ans.add(curAns);
            
            if (cur[2] < length2 - 1) {
                int index = cur[2] + 1;
                pq.offer(new int[]{cur[0], nums2[index], index});
            }
            k--;
        }
        
        return ans;
    }
}