力扣面试150 查找和最小的 K 对数字 最小堆 去重

Problem: 373. 查找和最小的 K 对数字

👨‍🏫 参考题解

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 创建一个大小为 k 的结果列表，用于存储和最小的 k 个数对
        List<List<Integer>> ans = new ArrayList<>(k); // 预分配空间
        
        // 创建一个优先队列（小根堆），存储三元组 [nums1[i] + nums2[j], i, j]
        // 按照和 (nums1[i] + nums2[j]) 的大小升序排列
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        
        // 将 nums1 中前 k 个元素与 nums2 中第一个元素的和及其索引 i, j 加入到优先队列中
        for (int i = 0; i < Math.min(nums1.length, k); i++) { // 至多 k 个
            pq.add(new int[]{nums1[i] + nums2[0], i, 0});
        }
        
        // 循环直到找到 k 个数对或者优先队列为空
        while (ans.size() < k && !pq.isEmpty()) {
            // 取出堆顶元素，也就是当前和最小的数对
            int[] p = pq.poll();
            int i = p[1]; // 取出 nums1 的索引
            int j = p[2]; // 取出 nums2 的索引
            
            // 将当前和最小的数对加入结果列表
            ans.add(List.of(nums1[i], nums2[j]));
            
            // 如果 nums2 中还有剩余元素，将新的数对 [nums1[i], nums2[j + 1]] 放入优先队列
            if (j + 1 < nums2.length) {
                pq.add(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
            }
        }
        
        // 返回结果列表
        return ans;
    }
}