LeetCode每日一题(373. Find K Pairs with Smallest Sums)

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u, v) which consists of one element from the first array and one element from the second array.

Return the k pairs (u1, v1), (u2, v2), …, (uk, vk) with the smallest sums.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]

Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]

Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]

Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Constraints:

• 1 <= nums1.length, nums2.length <= 105
• -109 <= nums1[i], nums2[i] <= 109
• nums1 and nums2 both are sorted in ascending order.
• 1 <= k <= 104

---

又是这种合并排序的问题， 我们还是需要建一个池子保存当前备选的选项， 这些选项由三个要素组成， index1 为 nums1 的 index， index2 为 nums2 的 index， sum 为 nums1[index1] + nums2[index2], 这三个当中 sum 是为排序准备的， 我们遍历 nums1 生成(nums1[i]+nums2[0], i, 0)放入一个 binary heap 中， 这样我们就得到了所有 nums1 中的元素与 nums2[0]所组合出来的备选项， 我们每次从中拿出最小的那个， 将 nums1[index1]和 nums2[index2]放到答案中， 然后我们将 index2+1 并重新计算备选项， 把备选项重新放到 binary heap 中就可以了。要注意的是 k 可能大于所有备选项的数量

---

use std::cmp::Reverse;
use std::collections::BinaryHeap;

impl Solution {
    pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, mut k: i32) -> Vec<Vec<i32>> {
        let mut heap = nums1
            .iter()
            .enumerate()
            .map(|(i, &n)| Reverse((n + nums2[0], i, 0usize)))
            .collect::<BinaryHeap<Reverse<(i32, usize, usize)>>>();
        let mut ans = Vec::new();
        while k > 0 {
            if let Some(Reverse((_, i1, mut i2))) = heap.pop() {
                ans.push(vec![nums1[i1], nums2[i2]]);
                k -= 1;
                i2 += 1;
                if i2 == nums2.len() {
                    continue;
                }
                heap.push(Reverse((nums1[i1] + nums2[i2], i1, i2)));
                continue;
            }
            break;
        }
        ans
    }
}