K Closest Points to Origin
Given an array of points on a plane and an integer k, return the k points closest to the origin.
Open on LeetCode ↗Intuition
Two traps, and the expensive one is silent. The first is calling sqrt on every distance — you are comparing distances, not reporting them, and sqrt is monotonic, so xx + yy orders the points identically while a float root only costs you precision and time. The second is heap direction. A min-heap of all n points feels right and gives the correct answer, but it holds every point at once: on a stream of ten million coordinates you have put ten million in memory to keep three. Flip it. Keep a MAX-heap capped at k, push each point, and the moment the size hits k+1 pop the root — in a max-heap that root is the farthest point you are holding, exactly the one you no longer want. O(n log k) time, O(k) space, never more than k points resident. The invariant is that after i points the heap holds precisely the k closest among them, with the worst of those at the root.
A bounded max-heap of size k: push, then evict the worst whenever the heap grows past k. What remains is the k best. Comparing squared distances avoids the square root entirely, since squaring preserves order on non-negative values.
Approach
Compare squared distance, never the root
The Euclidean distance from the origin is sqrt(xx + yy), but you are only ever asking which of two points is nearer. Since sqrt is strictly increasing on non-negative numbers, d1 < d2 if and only if sqrt(d1) < sqrt(d2) — the comparison is unchanged. Dropping the root removes a floating-point operation per point and, more importantly, removes floating-point rounding from a comparison that was exact integer arithmetic to begin with.
Keep a max-heap of size k, not a min-heap of size n
Python's heapq is a min-heap, so push the negated squared distance to simulate a max-heap. Push every point; whenever the heap exceeds k entries, pop the root. Because the root is the largest negated value, it is the smallest true distance's opposite — that is, the farthest point currently held. Evicting the farthest is always safe: if it is not among the k closest of what you have seen, it cannot become one of the k closest of anything larger.
Whatever survives is the answer, in any order
After the full pass the heap contains exactly k entries, and by the invariant those are the k closest points overall. The problem allows any order, so no final sort is required — just unpack the stored points. If you did want them ordered by distance you would pop them out, which adds O(k log k) and is the only place sorting would ever enter.
Solution & live demo
Common pitfalls
Taking the square root
d = math.sqrt(x * x + y * y)
d = x * x + y * y
The ordering by distance and by squared distance is identical, so the root is pure cost — and it introduces floating-point ties that can order equal-distance points inconsistently.
Using a min-heap and popping k times
heapq.heapify(all_points) return [heapq.heappop(heap) for _ in range(k)]
heapq.heappush(heap, (-d, x, y)) if len(heap) > k: heapq.heappop(heap)
That builds a heap of all n points — O(n) space and O(n + k log n) time. Capping at k keeps memory at O(k), which matters when n is large and k is small.
Evicting before pushing
if len(heap) == k: heapq.heappop(heap) heapq.heappush(heap, (-d, x, y))
heapq.heappush(heap, (-d, x, y)) if len(heap) > k: heapq.heappop(heap)
Popping first discards the current worst without comparing it to the incoming point, so a good point can be evicted in favour of a worse one. Push then evict guarantees the removed element is genuinely the furthest.
Edge cases
The heap never exceeds k, so nothing is ever evicted and every point is returned — correct by definition.
Its squared distance is 0, the smallest possible, so it is never the max-heap root while anything else is present and survives eviction.
Both have distance 1; the heap breaks the tie arbitrarily and either is a valid answer, since the problem guarantees the answer is unique only up to order.
Squaring stays exact in Python's unbounded integers, so no overflow — one of the concrete reasons to prefer xx + yy over a float sqrt.