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
from heapq import heapify, heappushpop


def k_closest(points, k, origin=(0, 0)):
    # 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.
    """
    heap = [(-distance(p, origin), p) for p in points[:k]]
    heapify(heap)

    """
    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.
    """
    for point in points[k:]:
        dist = distance(point, origin)

        heappushpop(heap, (-dist, point))  # 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 [point for nd, point in heap]  # return points in heap


def distance(point, origin=(0, 0)):
    """ Calculates the distance for a point from origo"""
    return (point[0] - origin[0])**2 + (point[1] - origin[1])**2
