Find k closest points in python

kclosest.py

import heapq


def k_closest_points(points_list, point, k):
    # Find k closest points to point in points_list
    points_distance = []
    for p in points_list:
        distance = (p[0] - point[0]) ** 2 + (p[1] - point[1]) ** 2
        points_distance.append((distance, p))
    heapq.heapify(points_distance)
    res = []
    for i in range(k):
        p = heapq.heappop(points_distance)
        res.append(p[1])
    return res

if __name__ == '__main__':
    assert k_closest_points([(1, 1), (2, 2), (3, 3), (4, 4), (0.5, 0.5)], (0, 0), 2) == [(0.5, 0.5), (1, 1)]