剑指offer 专项突破版 61、和最小的 k 个数对

题目链接

思路一:大根堆

• 既然要求最小的k的个数对，那我就维护一个大小为k的大根堆 遍历一遍，这样最后堆里所有的元素就是需要的结果了
• 注意并不需要全部遍历，因为只需要k个，所以只需要遍历每个数组的前k个元素就好了(但是又不知道k和数组长度哪个大，所以遍历次数是二者中小的那个)
• 注意如何设置优先队列是大根堆
• 时间复杂度：嵌套循环共遍历k²次，每次改变树的元素需要logk的时间，所以总共为O(k²logk)

class Solution {
    PriorityQueue<List<Integer>> priorityQueue;
    int k;

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        priorityQueue = new PriorityQueue<>((e1, e2) -> e2.get(0) + e2.get(1) - e1.get(0) - e1.get(1));
        this.k = k;

        for (int i = 0; i < Math.min(nums1.length,k); i++) {
            for (int j = 0; j < Math.min(nums2.length,k); j++) {
                add(nums1[i],nums2[j]);
            }
        }
        return priorityQueue.stream().toList();
    }

    private void add(int m, int n) {

        if (priorityQueue.size() < k)
            priorityQueue.add(List.of(m, n));
        else if (m + n < priorityQueue.peek().get(0) + priorityQueue.peek().get(1)) {
            priorityQueue.poll();
            priorityQueue.offer(List.of(m, n));
        }

    }
}

思路二：小根堆

• 这个思路十分巧妙：堆中存的是数对在数组中的索引，最开始把(0,0)(1,0)(2,0)······(k,0)添加到堆中(当然数目还是取Math.min(k,nums1.length)，然后循环k次或者是循环至堆空,每次都移除堆顶的元素(m,n)，然后把(m,n+1)放入堆
• 这种做法不必循环k²次，循环k次即可，把时间复杂度降低到了O(klogk)
• 中间很多判断都是因为可能 nums1.length * nums2.length < k

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        Queue<Integer[]> minHeap = new PriorityQueue<>((e1, e2) -> nums1[e1[0]] + nums2[e1[1]] - nums1[e2[0]] - nums2[e2[1]]);
        List<List<Integer>> result = new ArrayList<>();

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

        while (0 != k-- && !minHeap.isEmpty()) {
            Integer[] min = minHeap.poll();
            result.add(List.of(nums1[min[0]], nums2[min[1]]));
            if(min[1] + 1 < nums2.length)
                minHeap.offer(new Integer[]{min[0], min[1] + 1});
        }
        return result;
    }
}