K Closest Points to Origin (LeetCode 973)

The problem

You get a list of points [x, y] on a plane and an integer k. Return the k points closest to the origin (0, 0) by ordinary (Euclidean) distance. Any order is accepted.

Input:  points = [[3, 3], [5, -1], [-2, 4]], k = 2
Output: [[-2, 4], [3, 3]]

The squared distances are 18, 26 and 20. The two smallest are 18 ([3, 3]) and 20 ([-2, 4]).

How to think about it

Start with brute force. Sort all points by distance and take the first k. That’s O(n log n). It’s correct, but it orders all n points when you only need k of them.

First small insight: skip the square root. Distance is sqrt(x² + y²), and the square root keeps order (if a² < b², then a < b). So compare x² + y² directly. That’s faster and avoids floating-point issues.

The key insight: this is top-K, but for the smallest. To keep the k largest you use a min-heap and evict the smallest. To keep the k smallest distances, flip it: use a max-heap of size k and evict the farthest. Python’s heapq is a min-heap, so store the negated distance; then the heap’s root is the farthest point you’re holding, and popping it drops the worst candidate.

Which pattern is this? Top-K with a heap. The tell is “k closest”, “k nearest” or “k smallest” over a list too big to want to sort.

Approach

1. For each point, compute d = -(x² + y²).
2. Push (d, [x, y]) onto the heap.
3. If the heap has more than k items, pop the root: the farthest point held so far.
4. After the scan, the heap holds the k closest points. Return them (sorted here, just to make the output stable).

Solution

import heapq
def k_closest(points, k):
    h = []
    for x, y in points:
        d = -(x*x + y*y)                                     # negated squared distance: root = farthest
        heapq.heappush(h, (d, [x, y]))
        if len(h) > k: heapq.heappop(h)                      # evict the farthest of the k+1
    return sorted([p for _, p in h])
# k_closest([[1,3],[-2,2]], 1) -> [[-2,2]]

Dry run

On [[3, 3], [5, -1], [-2, 4]] with k = 2:

[5, -1] had the most negative d, meaning it was farthest, so it was evicted. The result is [[-2, 4], [3, 3]].

Complexity

• Time: O(n log k). Each point does a push and at most one pop on a heap of at most k + 1 items. The final sort is O(k log k), which is smaller.
• Space: O(k). The heap holds at most k + 1 points.

Edge cases and pitfalls

• k equals the number of points: nothing is ever popped and all points come back.
• Ties in distance (like [1, 0] and [0, 1]): any valid choice is accepted. The tuple then compares the point lists, which works in Python because lists are comparable.
• Forgetting to negate turns the heap into a “keep the k farthest” heap.
• Using math.sqrt is unnecessary and risks floating-point surprises.
• Faster on average: Quickselect on distance gives O(n) average time, at the cost of O(n²) worst case and more code.

Related problems

• Kth Largest Element in an Array (LeetCode 215): the same size-k trick with a min-heap.
• Top K Frequent Elements (LeetCode 347): top-K with a different key.
• Find K Pairs with Smallest Sums (LeetCode 373): k smallest again, but candidates are generated lazily.