数据结构 | 堆 | 问题3

问题描述

给定两个升序数组 nums1 和 nums2，以及两个整数 k 和 target，请你找到前 k 个元素对 (a, b)，满足 a 来自 nums1，b 来自 nums2，并且 a + b 唯一且最接近 target。

注意，被保证至少有一个对 (a, b) 满足 a + b = target，并且一定有 k 对不同的 (a, b) 满足 a + b 最接近 target。

数组中的每个元素只能被使用一次。

示例

以下是示例输入和输出：

输入：

nums1 = [1,7,11]
nums2 = [2,4,6]
k = 3
target = 16

输出：

[[5,11],[7,9],[6,10]]

解法

对于前面几篇堆问题的介绍，本题可以用一个小根堆来解决。

我们可以将数组 nums1 和数组 nums2 中的所有数相加，并将和与其对应的下标 $(i, j)$ 存入小根堆中。标记已经取到的元素可以用一个数组来实现。对于每次从堆中取出的最小元素 $(sum, i, j)$，我们都将其加入答案数组中，并且将标记数组中 $(i, j)$ 置为 True，表示这个元素已经被取走。

由于这个小根堆非常特殊，每次弹出的元素一定是堆中元素和最小的，同时也一定是最优的（因为如果这个元素不是最优的，那么必然存在另一个元素 $(sum', i', j')$，使得 $sum' < sum$，且 $(i', j')$ 的标记为 False，那么 $(sum', i', j')$ 一定在堆中，而 $(sum', i', j')$ 与 $(sum, i, j)$ 对于答案的贡献是一样的，因此如果取 $(sum', i', j')$ 作为答案而不是取 $(sum, i, j)$，并不会破坏答案的正确性；同时，若 $(i', j')$ 的标记为 True，那么它已经被取走过了，因此不可能还有 $(i', j')$ 配合其他元素一同成为最优解）。因此，我们可以保证小根堆弹出 $k$ 个元素后，得到的数组一定是所有可能的，且和最接近 target 的前 $k$ 个元素。

这个题需要注意的是，在添加元素到小根堆时，要先判断该元素对应的下标是否已经被标记，若已经标记，则说明该元素已经被取走，直接跳过即可。

代码实现

以下是使用 Python 实现的代码：

import heapq

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int, target: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)
        q, ans = [(nums1[0] + nums2[0], 0, 0)], []
        seen = set((0, 0))
        while k and q:
            _, i, j = heapq.heappop(q)
            ans.append([nums1[i], nums2[j]])
            k -= 1
            if i < m - 1 and (i + 1, j) not in seen:
                seen.add((i + 1, j))
                heapq.heappush(q, (nums1[i + 1] + nums2[j], i + 1, j))
            if j < n - 1 and (i, j + 1) not in seen:
                seen.add((i, j + 1))
                heapq.heappush(q, (nums1[i] + nums2[j + 1], i, j + 1))
        return ans

该算法的时间复杂度是 $O(k\log k)$，因为我们最多会从小根堆中取出 $k$ 个元素，每次弹出是 $O(\log k)$ 的，因此总时间复杂度是 $O(k\log k)$。空间复杂度是 $O(k)$，因为我们最终需要返回一个长度为 $k$ 的数组，堆中最多会有 $k$ 个元素。