k_near.py

from collections import namedtuple

Point = namedtuple("Point", "x y")

def euclidean_distance(p1, p2):
    return pow(p1.x - p2.x, 2) + pow(p1.y - p2.y, 2)

def near(points, k, n):
    def keyfunc(point):
        return euclidean_distance(k, point)
    return sorted(points, key=keyfunc)[:n]

if __name__ == '__main__':
    assert near([], Point(3, 5), 4) == []
    assert near([Point(10, 100), Point(3, 5), Point(4, 6)], Point(3, 5), 1) == [Point(3, 5)]
    assert near([Point(3, 5)], Point(3, 5), 10) == [Point(3, 5)]
    assert near([Point(10, 100), Point(3, 5), Point(4, 6)], Point(3, 5), 3) == [Point(3, 5), Point(4,6), Point(10, 100)]
    assert near([Point(10, 100), Point(3, 5), Point(300, 5000), Point(4, 6)], Point(3, 5), 3) == [Point(3, 5), Point(4,6), Point(10, 100)]