【LeetCode笔记 - 每日一题】373. 查找和最小的 K 对数字（Java、堆、优先队列）

题目描述

几天没打题，感觉脑子都是一团浆糊。。。。

• 升序：肯定得用这条件来优化复杂度
• 数对：用 int[2] 来表示
  

思路 && 代码

• 看了题解区彤哥的思路，代码其实就 8 行= =
• 堆初始化：按照[i, 0]的规律初始化堆
• 弹出：当前堆顶就是最小值，弹出[index1, index2]后，往后最有可能的数对是[index1, index2 + 1]，放入队列中继续循环～

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 小顶堆，每次弹出最小的当前结果
        PriorityQueue<int[]> heap = new PriorityQueue<>
            ((a, b) -> nums1[a[0]] + nums2[a[1]] - (nums1[b[0]] + nums2[b[1]])); // 策略模式

        // 堆初始化：把 nums1 的所有索引入队，nums2 的索引初始时都是 0
        // 优化：最多入队 k 个就可以了，因为提示中 k 的范围较小，这样可以提高效率
        for (int i = 0; i < Math.min(k, nums1.length); i++) heap.offer(new int[] {i, 0});

        List<List<Integer>> ans = new ArrayList<>();

        // 最多弹出 k 次（获取 k 对数字）
        while (k-- > 0 && !heap.isEmpty()) {
            // 弹出一个
            int[] pos = heap.poll();

            ans.add(Arrays.asList(nums1[pos[0]], nums2[pos[1]])); // 每次都弹出一个当前的最小值

            // 加入一个，将 index2 加 1 之后继续入队
            if (++pos[1] < nums2.length) heap.offer(pos);
        }

        return ans;
    }
}

• 无注释版

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> nums1[a[0]] + nums2[a[1]] - (nums1[b[0]] + nums2[b[1]])); 

        for (int i = 0; i < Math.min(k, nums1.length); i++) heap.offer(new int[] {i, 0});
        
        List<List<Integer>> ans = new ArrayList<>();

        while (k-- > 0 && !heap.isEmpty()) {
            int[] pos = heap.poll();
            ans.add(Arrays.asList(nums1[pos[0]], nums2[pos[1]])); 
            if (++pos[1] < nums2.length) heap.offer(pos);
        }
        return ans;
    }
}