【Lintcode】1274. Find K Pairs with Smallest Sums

题目地址：

https://www.lintcode.com/problem/find-k-pairs-with-smallest-sums/description

给定两个升序数组$A$和$B$，可以在两个数组分别取一个数求和。问和前$k$小的数对是哪些。要求结果有序，按照和小者优先排列，如果和一样，按照第一个数在$A$中的下标小者优先排序。如果总共和的个数不足$k$个，有多少个算多少个。

设$A=[a_0,a_1,...,a_{n-1}]$，$B=[b_0,b_1,...,b_{m-1}]$。将两者张成和矩阵：$$\left[\begin{array}{cccc}{a}_0+b_0& a_1+b_0& ...& a_{n-1}+b_0\\ a_0+b_1& a_1+b_1& ...& a_{n-1}+b_1\\ a_0+b_2& a_1+b_2& ...& a_{n-1}+b_2\\ ...& ...& ...& ...\\ a_0+b_{m-1}& a_1+b_{m-1}& ...& a_{n-1}+b_{m-1}\end{array}\right]$$所以，问题就转化为，在这个矩阵里求前$k$小。可以用最小堆来做。先将左上角入堆，然后每次出堆一个数对，就将其右边和下边的数对入堆。注意，这里需要判重，同一个数对不能重复入堆。这样，第$i$次出堆的就是第$i$小的和，加入最终答案即可。代码如下：

import java.util.*;

public class Solution {
    
    class Pair implements Comparable<Pair> {
        int idx1, idx2, sum;
        
        public Pair(int idx1, int idx2, int sum) {
            this.idx1 = idx1;
            this.idx2 = idx2;
            this.sum = sum;
        }
        
        @Override
        public boolean equals(Object o) {
            Pair pair = (Pair) o;
            return idx1 == pair.idx1 && idx2 == pair.idx2;
        }
        
        @Override
        public int hashCode() {
            return Objects.hash(idx1, idx2);
        }
    
        @Override
        public int compareTo(Pair o) {
            if (this.sum != o.sum) {
                return Integer.compare(this.sum, o.sum);
            }
            // 如果和相等，按照A中的下标排序
            return Integer.compare(this.idx1, o.idx1);
        }
    }
    
    /**
     * @param nums1: List[int]
     * @param nums2: List[int]
     * @param k:     an integer
     * @return: return List[List[int]]
     */
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // write your code here
        List<List<Integer>> res = new ArrayList<>();
        
        PriorityQueue<Pair> minHeap = new PriorityQueue<>();
        Pair start = new Pair(0, 0, nums1[0] + nums2[0]);
        minHeap.offer(start);
        
        Set<Pair> set = new HashSet<>();
        set.add(start);
        
        // 出堆k - 1次，就求出了前k - 1小的和的数对。
        // 由于有可能总共和就不足k - 1个，所以如果堆空了也要退出循环
        for (int i = 0; i < k - 1 && !minHeap.isEmpty(); i++) {
            Pair cur = minHeap.poll();
            res.add(Arrays.asList(nums1[cur.idx1], nums2[cur.idx2]));
            
            if (cur.idx1 < nums1.length - 1) {
                Pair next = new Pair(cur.idx1 + 1, cur.idx2, nums1[cur.idx1 + 1] + nums2[cur.idx2]);
                if (!set.contains(next)) {
                    minHeap.offer(next);
                    set.add(next);
                }
            }
            
            if (cur.idx2 < nums2.length - 1) {
                Pair next = new Pair(cur.idx1, cur.idx2 + 1, nums1[cur.idx1] + nums2[cur.idx2 + 1]);
                if (!set.contains(next)) {
                    minHeap.offer(next);
                    set.add(next);
                }
            }
        }
        
        // 如果堆不空，则堆顶就是第k小的和
        if (!minHeap.isEmpty()) {
            Pair cur = minHeap.peek();
            res.add(Arrays.asList(nums1[cur.idx1], nums2[cur.idx2]));
        }
        
        return res;
    }
}

时间复杂度$O(k\log k)$（出堆$k-1$次，入堆的数不足$2(k-1)$个数），空间$O(k)$。