K Closest Points to Origin

This is the mirror image of Kth Largest. You want the k smallest distances, so the heap you keep is a max-heap — and in Python that means negation. Most bugs in this problem are sign bugs.

The problem

You are given a list of points on a flat plane, each as [x, y], and a number k. Return the k points nearest to the origin (0, 0), measured by straight-line distance. Any order is fine, and the answer is guaranteed to be unique.

• points = [[3, 3], [5, -1], [-2, 4], [1, 1]], k = 2 → [[1, 1], [3, 3]]. The squared distances are 18, 26, 20 and 2, so the two smallest are 2 and 18.
• points = [[0, 1]], k = 1 → [[0, 1]].

Constraints: 1 ≤ k ≤ n ≤ 10⁴, coordinates between -10⁴ and 10⁴.

Clarifying questions

• Which distance? Straight-line (Euclidean): √(x² + y²).
• Order of the output? Any order.
• Ties at the boundary? Guaranteed not to happen.
• Duplicate points? Treat them as separate points.

Approach 1: sort by distance

def k_closest_sort(points: list[list[int]], k: int) -> list[list[int]]:    """Sort every point by squared distance and take the first k."""    return sorted(points, key=lambda p: p[0] * p[0] + p[1] * p[1])[:k]

O(n log n) time, O(n) space. It is short and correct, and at n = 10⁴ it is fast. The weakness is the same as before: it orders all n points to use k of them, and it needs every point in memory at once.

The key insight

Two separate ideas make this clean.

First, drop the square root. √ is monotonic: if a < b then √a < √b. So ordering points by x² + y² gives exactly the same order as ordering by the true distance. You avoid a slow function and, more importantly, you avoid floating-point numbers. Squared distances of integer points are exact integers.

Second, keep the k closest in a max-heap. The heap holds your best k candidates. When a new point arrives and the heap is too big, you must evict the worst candidate — the one farthest from the origin. The structure must surface the largest distance, which is a max-heap. In Python, push -distance so that the most negative number, which is the farthest point, sits at the root.

The rule is the same one as in Kth Largest, turned around. A small table is worth keeping in your head:

Approach 2: a size-k max-heap

1import heapq23def k_closest_heap(points: list[list[int]], k: int) -> list[list[int]]:4    """Keep the k closest points in a max-heap, built from negated distances."""5    heap: list[tuple[int, int, int]] = []    # (-distance, x, y): farthest on top6    for x, y in points:7        distance = x * x + y * y             # squared: no sqrt, same order8        heapq.heappush(heap, (-distance, x, y))9        if len(heap) > k:10            heapq.heappop(heap)              # evict the farthest of the k + 111    return [[x, y] for _, x, y in heap]

Dry run on [[3, 3], [5, -1], [-2, 4], [1, 1]] with k = 2. The table shows real distances; the heap stores them negated.

The function returns [[3, 3], [1, 1]] — the right two points, in heap order.

Complexity. Each point does one push and at most one pop on a heap of size k + 1: O(n log k) time, O(k) space. The tuple carries x and y after the key, and since they are integers, a tie on distance falls through to comparing integers, which is safe.

There is a second heap answer worth knowing, because candidates often write it first. Build a min-heap of all n points with heapify in O(n), then pop k times at O(log n) each. That is O(n + k log n) time, which is actually faster than O(n log k) when k is tiny. The cost is space: it holds all n points, and it cannot run on a stream. If the interviewer says "the points come from a sensor, forever", only the size-k max-heap still works. Name both, and pick by the constraints.

Approach 3: quickselect on distance

Like Kth Largest, this has an O(n) average solution. Partition the points around a random pivot distance until the k smallest sit in the first k positions. It uses O(1) extra space when done in place, but it is O(n²) in the worst case and cannot handle a stream. Mention it; write it only if asked.

Edge cases

• k = n: every point is returned; the heap never evicts.
• The origin itself [0, 0]: distance 0, always kept.
• Points with equal distance, like [1, 2] and [2, 1]: both have distance 5. The problem promises no tie at the boundary, so either both are in or both are out.
• Large coordinates: (10⁴)² + (10⁴)² is 2 × 10⁸, which fits easily; Python integers never overflow anyway. In Java or C++ this is still inside a 32-bit int, but only just — say so.

Follow-ups

• "Closest to a given point (a, b) instead of the origin." Use (x - a)² + (y - b)². Nothing else changes.
• "Points arrive as a stream." The heap version already handles it: keep the heap as state and return its contents on request.
• "Return them sorted from closest to farthest." Pop the heap into a list and reverse it, O(k log k) extra.