K Closest Points to Origin Solution

Problem

Given a list of points on a 2D plane and an integer k, return the k points nearest to the origin (0, 0). Distance is the standard Euclidean distance; you may return the k points in any order.

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: Squared distances are 18, 26, and 20 respectively — [3,3] and [-2,4] are the two closest.

Intuition

We need the k smallest distances out of n, but we don't need them sorted. That distinction matters: sorting all n points costs O(n log n) regardless of k, while a max-heap of size k costs only O(n log k) — a meaningful win when k is much smaller than n. The heap acts as a running "top-k" filter: evict the farthest point whenever a closer one arrives.

Approach 1 — Sort

Compute squared distance for every point, sort ascending by that key, and slice the first k.

Sort points in-place using squared distance as the key.
Return points[:k].

1from typing import List
2
3def kClosest(points: List[List[int]], k: int) -> List[List[int]]:
4    points.sort(key=lambda point: point[0]**2 + point[1]**2)
5    return points[:k]

Time: O(n log n) — dominated by sorting all n points.
Space: O(log n) — in-place sort stack depth.

Approach 2 — Max-Heap of Size k

Keep a max-heap bounded to k entries. Push each point; when the heap grows beyond k, pop the farthest. Whatever remains is the answer.

Initialize an empty max-heap.
For each point, compute its squared distance.
Push the point onto the heap.
If the heap now holds more than k entries, pop the maximum (farthest) entry.
Return all points remaining in the heap.

1import heapq
2from typing import List
3
4def kClosest(points: List[List[int]], k: int) -> List[List[int]]:
5    max_heap = []
6    for x, y in points:
7        dist_sq = x*x + y*y
8        # Negate dist_sq so Python's min-heap acts as a max-heap
9        heapq.heappush(max_heap, (-dist_sq, x, y))
10        if len(max_heap) > k:
11            heapq.heappop(max_heap)  # Evict the farthest point
12    return [[x, y] for _, x, y in max_heap]

Time: O(n log k) — each of the n points triggers at most one push and one pop on a heap of height log k.
Space: O(k) — the heap holds at most k + 1 entries at any moment.

Complexity Summary

Approach	Time	Space	When to use
Sort	O(n log n)	O(log n)	Simplest code; fine when k is close to n
Max-Heap	O(n log k)	O(k)	Best when k ≪ n, or when processing a stream of points

Common Mistakes

Computing square root unnecessarily — sqrt(x*x + y*y) adds float overhead and potential precision issues. Since sqrt is monotonically increasing, a² + b² < c² + d² always implies the same ordering as their square roots.
Negating distance in Python but forgetting to extract correctly — after storing (-dist_sq, x, y), the result must unpack as [[x, y] for _, x, y in max_heap], not [[-neg, x, y] ...]. Mixing up the sign produces garbage coordinates.
Building a min-heap over all n points and popping k times — this costs O(n + k log n) time and O(n) space, which is strictly worse than the max-heap approach for small k.
Assuming the sort approach preserves input order — points.sort(...) mutates the caller's list in-place. If the original ordering matters elsewhere, make a copy first.
Expecting sorted output — the problem asks for any k closest points; the heap result is in heap order, not distance order. Attempting to sort the output is wasted work.

K Closest Points to Origin — Solution