多路归并排序学习（leetcode 373）

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

给定两个以 升序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk) 。

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

提示:
1 <= nums1.length, nums2.length <= 105
-109 <= nums1[i], nums2[i] <= 109
nums1 和 nums2 均为升序排列
1 <= k <= 1000

思路一：单纯优先级队列算法，使用两个for循环将两个数组中的数加入优先级队列。当优先级队列的长度大于等于k时，比较堆顶数组值和新要加入的数组值，若堆顶数组值大于新要加入的数组值，则把堆顶数组出队，新要加入的数组值入队。

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>((o1,o2)->o2[0]+o2[1]-o1[0]-o1[1]);
        int rowN=Math.min(nums1.length,k);
        int colN=Math.min(nums2.length,k);
        for(int i = 0; i < rowN; i++){
            for(int j = 0; j < colN; j++){
                if(queue.size() < k){
                    queue.offer(new int[]{nums1[i],nums2[j]});
                }else if(queue.size()>=k){
                    int[] temp = queue.peek();
                    if(temp[0]+temp[1] > nums1[i] + nums2[j]){
                        queue.poll();
                        queue.offer(new int[]{nums1[i],nums2[j]});
                    }
                }
            }
        }
        List<List<Integer>> result = new ArrayList<>();
        while(!queue.isEmpty()){
            int[] temp2 = queue.poll();
            result.add(Arrays.asList(temp2[0],temp2[1]));
        }
        return result;
    }
}

思路二：多路归并排序，我们将这两个数组的首位元素下标以二元组 (i, j) 放入优先队列（小根堆），其中 i 为nums1[i] 的下标，j 为该点对中 nums2[j] 的下标，这步操作的复杂度为 O(nlogn)。始终确保nums1 为两数组中长度较少的那个，然后通过标识位来记录是否发生过交换，确保答案的点顺序的正确性。每次从优先队列（堆）中取出堆顶元素（含义为当前未被加入到答案的所有点对中的最小值），加入答案，并将该点对所在序列的下一位（如果有）加入优先队列中。

class Solution {
    boolean flag = true;
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> ans = new ArrayList<>();
        int n = nums1.length, m = nums2.length;
        if (n > m && !(flag = false)) 
            return kSmallestPairs(nums2, nums1, k);
        PriorityQueue<int[]> q = new PriorityQueue<>((a,b)->(nums1[a[0]]+nums2[a[1]])-(nums1[b[0]]+nums2[b[1]]));
        for (int i = 0; i < Math.min(n, k); i++) 
            q.add(new int[]{i, 0});
        while (ans.size() < k && !q.isEmpty()) {
            int[] poll = q.poll();
            int a = poll[0], b = poll[1];
            ans.add(new ArrayList<>(){{
                add(flag ? nums1[a] : nums2[b]);
                add(flag ? nums2[b] : nums1[a]);
            }});
            if (b + 1 < m) 
                q.add(new int[]{a, b + 1});
        }
        return ans;
    }
}