given a list of points find the k closest to the origin idea maintain a max heap of k elements we can iterate through all points if a point p has a smaller distance to the origin than the top element of a heap we add point p to the heap and remove the top element after iterating through all points our heap contains the k closest points to the origin time oknklogk space ok initialize max heap with first k points python does not support a max heap thus we can use the default min heap where the keys distance are negated for every point p in pointsk check if p is smaller than the root of the max heap if it is add p to heap and remove root reheapify same as if d heap00 heappushheap d p heappopheap note heappushpop is more efficient than separate push and pop calls each heappushpop call takes ologk time calculates the distance for a point from origo return point0 origin02 point1 origin12 time o k n k logk space o k initialize max heap with first k points python does not support a max heap thus we can use the default min heap where the keys distance are negated for every point p in points k check if p is smaller than the root of the max heap if it is add p to heap and remove root reheapify heappushpop does conditional check same as if d heap 0 0 heappush heap d p heappop heap note heappushpop is more efficient than separate push and pop calls each heappushpop call takes o logk time return points in heap calculates the distance for a point from origo
from heapq import heapify, heappushpop


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

    
    for point in points[k:]:
        dist = distance(point, origin)

        heappushpop(heap, (-dist, point))  
        

    return [point for nd, point in heap]  


def distance(point, origin=(0, 0)):
    
    return (point[0] - origin[0])**2 + (point[1] - origin[1])**2