LeetCode 373. Find K Pairs with Smallest Sums（最小堆）

373. Find K Pairs with Smallest Sums

Medium

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]

题意

给定两个升序数组nums1, nums2，定义一个“对”为在nums1和nums2中各取一个数得到的二元组，求k个和最小的“对”

思路

不妨设nums1的长度n1小于等于nums2的长度n2，将两个数组看成n1个有序数组：

• (nums1[0], nums2[0]), (nums1[0], nums2[1]), (nums1[0], nums2[2]), …, (nums1[0], nums2[n2-1])
• (nums1[1], nums2[0]), (nums1[1], nums2[1]), (nums1[1], nums2[2]), …, (nums1[1], nums2[n2-1])
• …
• (nums1[n1-1], nums2[0]), (nums1[n1-1], nums2[1]), (nums1[n1-1], nums2[2]), …, (nums1[n1-1], nums2[n2-1])
  从而转化为n1个有序数组merge问题，用最小堆来求解。
  一个经验是与“第k个”“前k个”有关的问题常常都是用最大堆/最小堆求解的。

代码

class Solution {
    class Pair implements Comparable<Pair> {
        int u, v, s;
        
        public Pair(int u, int v, int s) {
            this.u = u;
            this.v = v;
            this.s = s;
        }
        
        @Override
        public int compareTo(Pair o) {
            return s - o.s;
        }
    }
    
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int n1 = nums1.length, n2 = nums2.length;
        if (k == 0 || n1 == 0 || n2 == 0) {
            return Collections.emptyList();
        }
        if (k > n1 * n2) {
            k = n1 * n2;
        }
        boolean isSwap = false;
        if (n1 > n2) {
            int[] tmp = nums1;
            nums1 = nums2;
            nums2 = tmp;
            int tmpn = n1;
            n1 = n2;
            n2 = tmpn;
            isSwap = true;
        }
        ArrayList<List<Integer>> ans = new ArrayList<>();
        PriorityQueue<Pair> pq = new PriorityQueue<Pair>();
        int i = 0;
        for (i=0; i<n1; ++i) {
            pq.add(new Pair(i, 0, nums1[i] + nums2[0]));
        }
        for (i=0; i<k; ++i) {
            Pair pair = pq.poll();
            ArrayList<Integer> tmp = new ArrayList<>();
            if (!isSwap) {
                tmp.add(nums1[pair.u]);
                tmp.add(nums2[pair.v]);
            } else {
                tmp.add(nums2[pair.v]);
                tmp.add(nums1[pair.u]);
            }
            ans.add(tmp);
            if (pair.v + 1 < n2) {
                pq.add(new Pair(pair.u, pair.v + 1, nums1[pair.u] + nums2[pair.v+1]));
            }
        }
        return ans;
    }
}