2022/01/14 - 373. 查找和最小的 K 对数字

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

Description¶

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

给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1), (u2,v2) ... (uk,vk) 。

就是给你两个数组，找出和最小的前k个数对，例如：

nums1 = [1, 2, 3, 4, 5], nums2 = [2, 5, 8, 9, 10], k = 2 那么答案就是：[[1, 2], [2, 2]]

Solution¶

通过题目描述，我们是可以大概确定解题方法了，这道题类似于 TopK 问题，但是这里我们有两个序列，而且是对数对的排序。首先能想到的就是遍历所有的数对，然后取前k个小数对。但是题目规定nums1和nums2的长度最大都为 10^5，遍历所有的数对的话，时间复杂度太高，会出现TLE。因为我们只关心 Top K 个数对，而两个数组又是升序的，所以我们能先将前K个数对放入小根堆，取出堆顶数对，再放入下一个数对。

主要思想是多路归并，首先[nums1[0], nums2[0]]肯定是最小的数对，那么下一个小的数对肯定就是在(nums1[0], nums2[1])和(nums1[1], nums2[0])之间，也就是说对于某个索引i和j的数对，下一个比它大的最小数对肯定是在(nums1[i], nums2[j + 1])和(nums1[i + 1], nums2[j])之间，那我们的思路差不多就清晰了。

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> ans = new ArrayList<>();
        int n = nums1.length, m = nums2.length;
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (a, b) -> (nums1[a[0]  ] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]]));
        for (int i = 0; i < Math.min(n,  k); i++) {
            pq.offer(new int[] {i, 0});
        }
        while (!pq.isEmpty() && k > 0) {
            int[] pair = pq.poll();
            ans.add(List.of(nums1[pair[0]], nums2[pair[1]]));
            if (pair[1] + 1 < m) pq.offer(new int[] {pair[0], pair[1] + 1});
            k--;
        }
        return ans;
    }
}

2022/01/14 - 373. 查找和最小的 K 对数字

Description¶

Solution¶

评论