k最接近原点

Question

我有一个点列表：[(x, y), (x, y), (x ,y) ... (x, y)]。 我希望k最近的点到(0, 0)。

我正在尝试实现this链接中的内容。但是，我正在错误地实现算法，我不确定它出错的地方。我想也许heapify不知道如何维持点之间的顺序。我怎么解决这个问题？

码

import matplotlib.pyplot as plt
from random import randint
from heapq import heappush, heappop, heapify
from math import sqrt

def distance(pointA, pointB):
    return sqrt((pointB[0] - pointA[0]) ** 2 + (pointB[1] - pointA[1]) ** 2)

def closest(points, k, origin):
    heap = []
    for point in points[:k]:
        heappush(heap, point)

    for point in points[k:]:
        if distance(point, origin) < distance(heap[0], origin):
            heappop(heap)
            heappush(heap, point)
    return heap

def naive(points, k, origin):
    sortedPoints = sorted(points, key=lambda p: distance(p, origin))
    return sortedPoints[:k]

points = [(randint(0, 100), randint(0, 100)) for i in range(100)]
k = 4
resA = closest(points, k, (0, 0))
resB = naive(points, k, (0, 0))
plt.scatter(*zip(*points))
plt.scatter(*zip(*resA))
plt.scatter(*zip(*resB))
plt.show()

结果

绿色的点由朴素方法给出，橙色的点由方法使用堆给出。

Answer 1

解决方案中的堆不变量使用该点的第一个元素。您想使用原点的距离：

def closest(points, k, origin):

    heap = [(-distance(p, origin), p) for p in points[:k]]
    heapify(heap)

    for p in points[k:]:
        d = distance(p, origin)
        heappushpop(heap, (-d, p))
    return [p for nd, p in heap]

注意：我还从heapq导入heappushpop，因为它比单独的来电更有效。

修改：删除了对heappushpop调用的条件环绕，因为条件也存在于该函数内。

Answer 2

由于您已经在使用heapq，因此您也可以使用函数而不是重新创建它：

from heapq import nsmallest
result = nsmallest(k, points, lambda p : distance(p, (0, 0)))