原贴地址:https://dawnmagnet.github.io/algorithm-station/docs/2022/1/14
给定两个以 升序排列 的整数数组 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]
示例 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] <= 10E9nums1 和 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个元素。
以上就是主要的思路,接下来讲一些可以优化的点(不优化可能会超时):
- 因为我们来看到题目的数据范围,
nums1和nums2中都有可能存在至多$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
}
}