K Closest Points to Origin

Q: 距离原点最近的 K 个点的时间复杂度是多少？

对所有点进行排序需要 O(n log n) 。堆版本的扫描需要 O(n log k) ，之后对剩余的 k 个点排序需要 O(k log k) ，并且内存中只保留 O(k) 个点。

Q: 为什么要使用最大堆，而不是最小堆，来查找距离最近的 k 个点？

一个新点会与已保留的 k 个点中距离最远的点竞争，因此这个点必须位于堆顶。只有当最小堆保存所有 n 个点时，它才能移除距离最近的 k 个点，但这样需要 O(n) 的内存，而不是 O(k) 。

Q: 比较距离时需要用到平方根吗？

不。对于非负数，只有当 a < b 时， √a < √b 才成立，因此 x² + y² 得出的顺序与距离相同。它也仍然是整数：这里最大为 2 × 10^8 ，在 32 位范围内。

Q: 能否在 O(n) 时间内解决 K 个离原点最近的点问题？

平均来说，是的。Quickselect 会按照相同的 (x² + y², x, y) 顺序围绕一个随机基准点对这些点进行分区，并且只在包含位置 k 的一侧继续操作，平均时间复杂度为 O(n) 。位于该位置左侧的 k 个点就是答案；对它们排序会增加 O(k log k) 的开销。

中等堆排序

+10

给定平面上的 n 个点：点 i 是 (xs[i], ys[i])。返回距离原点 (0, 0) 最近的 k 个点，每个点表示为一对 [x, y]。

按点到原点的距离从近到远排列答案。距离相同的点先按 x 从小到大排列，再按 y 从小到大排列。重复列出的点按每个副本分别计数。

函数

kClosest(xs: integer-array, ys: integer-array, k: integer) → integer-2d-array

xsinteger-array: 每个点的 x 坐标
ysinteger-array: 每个点的 y 坐标
kinteger: 返回多少个点
返回值integer-2d-array: 按距离从近到远排列的 k 个最近点，以 [x, y] 对的形式表示；距离相同时，按 x 再按 y 排序

约束条件

n == xs.length == ys.length
1 ≤ k ≤ n ≤ 5000
-10⁴ ≤ xs[i], ys[i] ≤ 10⁴
积分可能会重复，每个副本都是一个单独的积分。

示例

输入: xs = [3, -1, 2, 0]ys = [1, 2, -2, 5]k = 2
输出: [[-1, 2], [2, -2]]
解释: 平方距离分别是 10、5、8 和 25。两个最小值依次对应 (-1, 2)（5）和 (2, -2)（8）。

输入: xs = [1, -1, 0, 1]ys = [1, 1, 2, -1]k = 3
输出: [[-1, 1], [1, -1], [1, 1]]
解释: (1, 1)、(-1, 1) 和 (1, -1) 的平方距离都是 2，而 (0, 2) 的平方距离是 4。在距离相同的点中，(-1, 1) 的 x 最小；而 (1, -1) 按 y 排在 (1, 1) 前面。

输入: xs = [7]ys = [-7]k = 1
输出: [[7, -7]]
解释: 只有一个点，而你需要一个点，所以答案就是 (7, -7)。

进阶

现在这些点会逐个到达，每到达一个点，你都必须报告目前为止最近的 k 个点。你需要多少内存？

提示

逐个打开提示，每个都会多透露一点。

已打开 0/3

题解

距离计算不需要开平方根：x² + y² 得到的排序相同，而且结果仍是整数。你只需要 n 个点中的 k 个，因此，用一个最大堆保存目前最近的 k 个点，并让最远的点位于堆顶，以便随时替换，时间复杂度为 O(n log k)，而不是对所有点进行完整排序。

选择距离最近的点 k 次

结果正确，但无法在最大的测试中运行完毕

思路

答案中的第一个点是所有点中最近的，第二个点是剩余点中最近的，以此类推。因此，扫描这些点 k 次，每次从尚未选取的点中，按照 (x² + y², x, y) 选出最靠前的点。

每次扫描都会遍历全部 n 个点。对于 n = 5000 和 k = 2500，这需要进行 1.25 × 10^7 次三部分比较。

算法

将每个点标记为未选取。
重复 k 次：扫描未选取的点，找出按 (x² + y², x, y) 排序后最靠前的点。
将它标记为已选取，并将 [x, y] 添加到答案中。
返回答案。

def kClosest(xs, ys, k):
    n = len(xs)
    taken = [False] * n
    result = []
    for _ in range(k):
        best = -1
        for i in range(n):
            # Compare (squared distance, x, y): the order the answer uses.
            if not taken[i] and (best == -1 or
                                 (xs[i] * xs[i] + ys[i] * ys[i], xs[i], ys[i]) <
                                 (xs[best] * xs[best] + ys[best] * ys[best], xs[best], ys[best])):
                best = i
        taken[best] = True
        result.append([xs[best], ys[best]])
    return result

对每个点进行排序

思路

答案是完整排序中前 k 个点，因此按 (x² + y², x, y) 对所有点排序，并保留前 k 个。在示例 1 中，排序后的键是 (5, -1, 2)、(8, 2, -2)、(10, 3, 1) 和 (25, 0, 5)，前两个就是答案。

排序的时间开销为 O(n log n)，与 k 的取值无关。当 k 很小时，大部分排序工作都花在了那些最终不会进入答案的点上。

算法

构建每个点的键 (x² + y², x, y)。
按键对点进行排序。
将前 k 个点作为 [x, y] 对返回。

def kClosest(xs, ys, k):
    points = list(zip(xs, ys))
    # Sort by squared distance, then x, then y: the order the answer uses.
    points.sort(key=lambda p: (p[0] * p[0] + p[1] * p[1], p[0], p[1]))
    return [[x, y] for x, y in points[:k]]

在最大堆中保留距离最近的 k 个

思路

保留目前为止见过的最优 k 个点。新点只能挤掉保留点中排名最后的那个，因此用最大堆按 (x² + y², x, y) 存储保留点，并让排名最后的点位于堆顶。存入 k 个点后，只有排名靠前的点才能替换堆顶；其他点都会被丢弃，因为已有 k 个点排名都在它前面。

在示例 2 中，当 k = 3 时，存入前三个点后，堆顶是 (0, 2)，其键为 (4, 0, 2)。点 (1, -1) 的键为 (2, 1, -1)，排名靠前，因此会替换它。

堆只会部分排序，因此移除堆顶 k 次，并从答案的最后一个位置开始依次填入，直到第一个位置。每个点最多需要进行一次 O(log k) 的堆操作，因此时间复杂度为 O(n log k)，空间复杂度为 O(k)。

算法

创建一个以 (x² + y², x, y) 为键的空最大堆。
插入前 k 个点。
对于后续的每个点，将它与堆顶进行比较：如果它的排名更靠前，就将它换入并向下筛；否则跳过它。
移除堆顶 k 次；每次移除的点都会填入答案中最后一个空位。
返回答案。

import heapq


def kClosest(xs, ys, k):
    # heapq is a min-heap; negating every part of (distance, x, y) puts the
    # farthest kept point, in the answer's order, on top.
    heap = []
    for x, y in zip(xs, ys):
        item = (-(x * x + y * y), -x, -y)
        if len(heap) < k:
            heapq.heappush(heap, item)
        elif item > heap[0]:
            # The new point comes before the farthest kept one: swap them.
            heapq.heapreplace(heap, item)
    # Popping gives the farthest first, so fill the answer from the back.
    result = [None] * len(heap)
    for i in range(len(result) - 1, -1, -1):
        _, x, y = heapq.heappop(heap)
        result[i] = [-x, -y]
    return result

易错点和边界情况

跳过平局时的比较。例如 2 中有 3 个点的距离都是 2；如果不先比较 x，再比较 y，它们就会按输入顺序排列，而不是按 (-1, 1)、(1, -1)、(1, 1) 的顺序排列。
维护一个包含 k 个点的最小堆。堆顶是最近的点，因此每次移除都会丢弃答案中需要的点，最后留下的会是距离最远的 k 个点。
直接从堆的数组中读取答案。这个数组只经过部分排序；你应该先逐个移除这些点。
移除重复的点。如果 (5, 5) 列出了 3 次，且 k = 3，那么答案就是这个点 3 次。

常见问题4

距离原点最近的 K 个点的时间复杂度是多少？

对所有点进行排序需要 O(n log n)。堆版本的扫描需要 O(n log k)，之后对剩余的 k 个点排序需要 O(k log k)，并且内存中只保留 O(k) 个点。

为什么要使用最大堆，而不是最小堆，来查找距离最近的 k 个点？

一个新点会与已保留的 k 个点中距离最远的点竞争，因此这个点必须位于堆顶。只有当最小堆保存所有 n 个点时，它才能移除距离最近的 k 个点，但这样需要 O(n) 的内存，而不是 O(k)。

比较距离时需要用到平方根吗？

不。对于非负数，只有当 a < b 时，√a < √b 才成立，因此 x² + y² 得出的顺序与距离相同。它也仍然是整数：这里最大为 2 × 10^8，在 32 位范围内。

能否在 O(n) 时间内解决 K 个离原点最近的点问题？

平均来说，是的。Quickselect 会按照相同的 (x² + y², x, y) 顺序围绕一个随机基准点对这些点进行分区，并且只在包含位置 k 的一侧继续操作，平均时间复杂度为 O(n)。位于该位置左侧的 k 个点就是答案；对它们排序会增加 O(k log k) 的开销。

def kClosest(xs, ys, k):
    # 在这里编写代码

测试 #1

测试 #2

测试 #3

参数

xs = [3, -1, 2, 0]
ys = [1, 2, -2, 5]
k = 2

预期输出

[[-1, 2], [2, -2]]