算法-堆/多路归并-查找和最小的 K 对数字

news2024/11/15 6:33:21

算法-堆/多路归并-查找和最小的 K 对数字

1 题目概述

1.1 题目出处

https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/description/?envType=study-plan-v2&envId=top-interview-150

1.2 题目描述

在这里插入图片描述

2 优先级队列构建大顶堆

2.1 思路

将两个数字的和放入大顶堆中，堆的最大大小为k。

当堆大小小于k时，直接放里面放。
当堆大小达到k后，比较当前元素和堆顶的元素，如果比堆顶元素小，就移除堆顶元素并放入当前元素。

最后，堆内元素就是和最小的K对数。

2.2 代码

class Solution {
    List<List<Integer>> resultList = new LinkedList<>();
    

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 大顶堆
        PriorityQueue<List<Integer>> queue = new PriorityQueue<>((o1, o2)->o2.get(0) + o2.get(1) - o1.get(0) - o1.get(1));
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            for (int j = 0; j < Math.min(nums2.length, k); j++) {
                int sum = nums1[i] + nums2[j];
                if (queue.size() < k) {
                    List<Integer> pair = new ArrayList<>();
                    pair.add(nums1[i]);
                    pair.add(nums2[j]);
                    queue.add(pair);
                } else {
                    List<Integer> headPair = queue.peek();
                    int headSum = headPair.get(0) + headPair.get(1);
                    if (sum < headSum) {
                        queue.poll();
                        List<Integer> pair = new ArrayList<>();
                        pair.add(nums1[i]);
                        pair.add(nums2[j]);
                        queue.add(pair);
                    }
                }
            }
        }
        while (queue.size() > 0) {
            resultList.add(queue.poll());
        }
        return resultList;
    }
}