k 离原点最近的点

Question

我正在尝试解决这个问题。 https://leetcode.com/problems/k-closest-points-to-origin/ 我使用过 maxheap，下面是我的代码。这适用于 points= [[1,3],[-2,2]], k = 1。但是对于 points = [[3,3],[5,-1],[-2,4]] , k = 2, 这给出的输出为 [[3,3],[5,-1]] 而预期输出为 [[3,3],[-2,4]]。我无法弄清楚为什么我的代码选择点 [5,-1]。请帮忙找出我犯的错误。

from heapq import *
class Solution:

def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:

    maxHeap = []
    for i in range(k):
        heappush(maxHeap, points[i])
        
    for i in range(k, len(points)):
        dist = self.distance(points[i])
        heapdistance_from_origin = self.distance(maxHeap[0])
        if dist < heapdistance_from_origin:
            heappop(maxHeap)
            heappush(maxHeap, points[i])
    return maxHeap

def distance(self, point: List[int]):
     return point[0] * point[0] + point[1] * point[1]

Answer 1

问题： heappush 和 heappop 不按到原点的距离对堆元素进行排序。他们改用元素的自然顺序（它们是列表，所以它是字典顺序）。所以 maxHeap[0] 通常 而不是 离原点最远的点。

您可以使用 heapq.nsmallest，它支持 key 参数（并且它使用与您几乎相同的算法）：

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return nsmallest(k, points, key=lambda p: hypot(*p))

或者：

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return nsmallest(k, points, key=partial(dist, (0, 0)))

顺便说一句，不需要自己导入，LeetCode 已经为你导入了几乎所有的东西。（当然，如果您想要清楚起见，您仍然可以这样做，那么在我的例子中，很明显 hypot 和 dist 来自 math 而 partial来自 functools。LeetCode 应该是关于编码面试的，我猜在白板面试中，大声说出“partial from functools”很方便，写得更少，也许这就是 LeetCode 这样做的原因。 )

如果你想自己实现它： 你可以让它保存成对的 (negated_distance, point) 而不是让你的堆保存裸点。取反是因为 heapq 实现了 min-heap，因此如果您想要 max-heap，则取反这些值。这是：

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:

        maxHeap = []
        for p in points[:k]:
            heappush(maxHeap, (-self.distance(p), p))

        for p in points[k:]:
            dist = self.distance(p)
            heapdistance_from_origin = -maxHeap[0][0]
            if dist < heapdistance_from_origin:
                heappop(maxHeap)
                heappush(maxHeap, (-dist, p))
        return [p for _, p in maxHeap]

或者您可以将点转换为 Point class 的实例，以进行您想要的比较：

    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:

        class Point(list):
            def __lt__(self_, other):
                return self.distance(self_) > self.distance(other)
        points = list(map(Point, points))
        
        maxHeap = []
        for i in range(k):
            heappush(maxHeap, points[i])

        for i in range(k, len(points)):
            dist = self.distance(points[i])
            heapdistance_from_origin = self.distance(maxHeap[0])
            if dist < heapdistance_from_origin:
                heappop(maxHeap)
                heappush(maxHeap, points[i])
        return maxHeap

k 离原点最近的点

k closest point to origin

python

max-heap