SWE · Onsite Coding

Apple 有序数组 Top K 最近数对:堆归并所有行

按固定左端点把全部数对拆为单调行,用小顶堆得到全局最近的 K 对。

题目重点:按固定左端点把全部数对拆为单调行,用小顶堆得到全局最近的 K 对。

题意与假设

这道练习题保留了面经中值得训练的核心约束,并将需要澄清的条件单独列出。现场拿到变体时,先和面试官确认这些条件。

  • 输入非降序,数对由两个不同位置 i<j 构成;值相同的不同位置仍是不同候选数对。
  • 按 (距离,i,j) 排序打破并列;输出值对,若要区分重复值位置应输出索引。k<0 报错,k 大于全部数对数则返回全部数对。
  • 这里是所有 n(n−1)/2 个位置对,不能只在 n−1 个相邻位置对中挑 K 个。

解题路线

  1. 固定左端点 i,候选 (i,i+1),(i,i+2),... 的距离 `a[j]−a[i]` 非降序,形成一条已排序行。
  2. 把每行第一个候选 `(distance,i,j)` 放入小顶堆;每弹出一对,只把同一行的下一个 j+1 入堆。
  3. 堆始终含每条非空行尚未输出的最小数对;全局最小必在这些行头中,归纳保证每次弹出的是全局下一个。
  4. 与显式生成并按 `(距离,i,j)` 排序所有数对的独立参考实现对照 30,459 个数组/K 组合。

手动走一遍

数组 [1,2,4,7,11,16]:首批行头距离为 1、2、3、4、5;弹 (1,2) 后推同一行 (1,4) 距 3,下一次仍弹 (2,4) 距 2。若 K 扩大,(1,4) 这种非相邻元素也可能进入结果。

Python 参考实现

下面的代码按上文假设实现,适合本地运行并追加自己的测试用例。

代码区可横向滚动;使用键盘时,先聚焦代码区再按方向键查看。

from heapq import heapify, heappop, heappush


def top_k_closest_pairs(values: list[int], k: int) -> list[tuple[int, int]]:
    """Sorted input; all i<j pairs ranked by (distance, i, j)."""
    if k < 0:
        raise ValueError("k must be non-negative")
    if any(values[i] > values[i + 1] for i in range(len(values) - 1)):
        raise ValueError("input must be sorted")
    n = len(values)
    heap = [(values[i + 1] - values[i], i, i + 1) for i in range(n - 1)]
    heapify(heap)

    answer = []
    while heap and len(answer) < k:
        _, i, j = heappop(heap)
        answer.append((values[i], values[j]))
        if j + 1 < n:
            heappush(heap, (values[j + 1] - values[i], i, j + 1))
    return answer

复杂度

建堆 O(n),每输出一对至多一次弹出和一次插入,O(k log n);总时间 O(n+k log n)、堆空间 O(n)、输出 O(min(k,n(n−1)/2))。若 K 接近全部对数,完整枚举排序可作为简单基线。

易错点与追问

容易失分的边界

  • 只计算相邻差,K 大时遗漏非相邻数对。
  • 把重复值的不同位置折叠,丢失合法零距离数对。
  • 未定义同距离的稳定顺序,使输出和测试不一致。
  • 输入并非已排序却直接按行单调性推进。
  • K 大于总对数时继续弹空堆。

继续深挖

  • 如果输入无序,先排序后返回原始位置如何做?
  • 若要求只返回唯一值对,如何避免重复输出?
  • 若距离是二维欧氏距离,固定 i 的行还单调吗?
  • 如果 k 非常小,能否减少初始化 n−1 行的开销?
✓已复制微信号!