跳过正文

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

·1249 字·3 分钟
曙光磁铁
作者
曙光磁铁

原贴地址:https://dawnmagnet.github.io/algorithm-station/docs/2022/1/14

给定两个以 升序排列 的整数数组 nums1nums2,以及一个整数 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]

示例 2:

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

示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

提示:

  • 1 <= nums1.length, nums2.length <= 10E5
  • -10E9 <= nums1[i], nums2[i] <= 10E9
  • nums1 和 nums2 均为升序排列
  • 1 <= k <= 1000

来源:力扣(LeetCode) 链接:https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums 著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。

思路分析
#

本题主要思路就是一个找最小k个,一般找最小k个或者最大k个都会用到一个算法叫做堆,这个堆一般也有很多叫法,有叫最小堆、最大堆的,也有叫优先队列的,但本质都是维护一个堆。这个堆有一些性质,最大堆的每个父节点都比子节点大,最小堆每个父节点都比子节点小。在维护这个性质的同时,我们就可以对其进行入堆或者出堆操作,并用一个非常短的时间$O(lg(n))$来维护这个堆的性质。这样就能做到快速获得几个最大或者最小的元素。

在这里,因为我们总体要来查询的一共有$m*n$个元素,也就是两个数组两两相加的数字都有可能是结果的一部分。

所以我们维护一个最大堆,在这里我们就叫它heap,这个heap存在的元素个数始终不会超过k个,k也就是最终答案里所要求的答案个数,首先我们来两层循环,也就是两两相加,往这个heap里放置元素。因为我们最终要保留k个,所以一旦heap里元素的个数大于k个,就将其pop掉。在这种情况下我们仍然可以保证一定可以取得到最小的k个元素。

以上就是主要的思路,接下来讲一些可以优化的点(不优化可能会超时):

  • 因为我们来看到题目的数据范围,nums1nums2中都有可能存在至多$10^5$个元素,而我们最终的答案最多最多也就1000个,都判断肯定没有必要,所以首先可以限制循环的变量变化的范围,原来是不设限制,现在可以改成最大到k(因为最终的答案一定是在nums1中前k个数字和nums2中前k个数字两两配对的结果中的,稍微仔细想一想就能明白吧)
  • 第二个优化的地方就是在循环过程中,因为我们是按顺序从前向后遍历的,如果当前遍历到的数字大于等于堆中的最大元素时,那么对于当前的某些循环已经没有遍历的必要了,就可以进行break跳过。

Rust代码
#

use std::collections::BinaryHeap;
impl Solution {
    pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<Vec<i32>> {
        let mut res = vec![];
        let mut heap = BinaryHeap::new();
        for i in 0..(nums1.len()).min(k as usize) {
            for j in 0..(nums2.len()).min(k as usize) {
                heap.push((nums1[i] + nums2[j], nums1[i], nums2[j]));
                if heap.len() > k as usize {
                    if let Some((a, b, c)) = heap.pop() {
                        if b == nums1[i] && c == nums2[j] {
                            break;
                        }
                    }
                }
            }
        }
        
        while heap.len() > 0 {
            if let Some((a, b, c)) = heap.pop() {
                res.push(vec![b, c]);
            }
        }
        res
    }
}