K Closest Points to Origin

1. The obvious way

   Sort all points by distance and take the first k: O(n log n).

2. The key insight

   Keep a max-heap (farthest on top) of at most k points. Push each point with its squared distance; if the heap has more than k, pop the farthest. At the end the heap holds the k closest. O(n log k).

3. The algorithm
   1. For each point compute distance squared = x*x + y*y.
   2. Push it into a max-heap keyed by that distance.
   3. If the heap has more than k, pop the farthest.
   4. Return the points left in the heap.

Watch it run

heap: 18, 26

import heapq

def k_closest(points, k):
    heap = []                                          # max-heap by distance (negatives)
    for x, y in points:
        distance = x * x + y * y                       # no square root needed for ordering
        heapq.heappush(heap, (-distance, x, y))
        if len(heap) > k:
            heapq.heappop(heap)                        # drop the farthest
    return [[x, y] for _, x, y in heap]

Now close this and write it again from memory, in the language you will use in the interview. Reattempt after 1 day, 7 days and about 30 days.