Top-K / Heap — K Closest Points to Origin
↗ LeetCode: K Closest Points to Origin · Medium · pattern chapter →
Given points[][2] and integer k, return the k points closest to origin (0, 0) (Euclidean).
Example 1 — points = [[1,3],[-2,2]], k = 1 → [[-2,2]]Example 2 — points = [[3,3],[5,-1],[-2,4]], k = 2 → [[3,3],[-2,4]]Example 3 — points = [[0,0],[1,1]], k = 1 → [[0,0]]
Constraints — 1 ≤ k ≤ n ≤ 10⁴; -10⁴ ≤ x, y ≤ 10⁴. Any order accepted. Brute sorts all n by distance — O(n log n) = 10⁴·14 ≈ 10⁵ ops at n=10⁴, ~10⁶ ops streaming 10 requests/sec for 1 min. Max-heap of size k is O(n log k) = 10⁴ ops per call; quickselect is expected O(n) = 10⁴ ops.
Interview timer00:00
Approach 1 — Sort by distance
Intuition. Sort all points by squared distance; take first k.
Trap — use squared distance x² + y², not sqrt (avoid FP and unnecessary cost).
int[][] kClosestSort(int[][] points, int k) {
Arrays.sort(points, (a, b) -> (a[0]*a[0] + a[1]*a[1]) - (b[0]*b[0] + b[1]*b[1]));
return Arrays.copyOfRange(points, 0, k);
}Complexity — Time O(n log n); Space O(1) in-place sort.
Approach 2 — Max-heap of size k
Insight from sort. We don't need to sort everything — maintain a max-heap of size k on distance. Push each point; if size exceeds k, pop the farthest. The heap always holds the k closest seen so far.
int[][] kClosest(int[][] points, int k) {
PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) ->
(b[0]*b[0] + b[1]*b[1]) - (a[0]*a[0] + a[1]*a[1]));
for (int[] p : points) {
pq.offer(p);
if (pq.size() > k) pq.poll();
}
return pq.toArray(new int[0][]);
}Complexity — Time O(n log k); Space O(k).
Approach 3 — Quickselect (avg O(n))
Insight from heap. Since we don't need the k closest in order, we can partition around the k-th smallest — like Quickselect. Pick pivot; partition; recurse on the side containing the boundary.
int[][] kClosestQS(int[][] points, int k) {
quickselect(points, 0, points.length - 1, k);
return Arrays.copyOfRange(points, 0, k);
}
void quickselect(int[][] a, int lo, int hi, int k) {
if (lo >= hi) return;
int p = partition(a, lo, hi);
if (p == k) return;
if (p < k) quickselect(a, p + 1, hi, k);
else quickselect(a, lo, p - 1, k);
}
int partition(int[][] a, int lo, int hi) {
int pivotDist = dist(a[hi]);
int i = lo;
for (int j = lo; j < hi; j++)
if (dist(a[j]) < pivotDist) swap(a, i++, j);
swap(a, i, hi);
return i;
}
int dist(int[] p) { return p[0]*p[0] + p[1]*p[1]; }
void swap(int[][] a, int i, int j) { int[] t = a[i]; a[i] = a[j]; a[j] = t; }Complexity — Time O(n) average; O(n²) worst; Space O(1) in place (recursion O(log n) avg).
Try it yourself
Complexity summary
| Approach | Time | Space | Interview grade |
|---|---|---|---|
| Sort by distance | O(n log n) | O(1) | baseline; simplest |
| Max-heap of size k | O(n log k) | O(k) | canonical Top-K answer |
| Quickselect | O(n) avg | O(1) | best asymptotic when order doesn't matter |
When to use which
- Order matters within top-k → max-heap or sort.
- Order doesn't matter, best avg time → Quickselect.
- Streaming (points arrive one by one) → max-heap.
- k close to n → sort or min-heap of size (n-k).
Related problems
- Top K Frequent Elements — canonical Top-K
- Kth Largest Element in an Array — Quickselect archetype
- Kth Largest Element in a Stream — streaming
- Kth Smallest Element in a Sorted Matrix — BS-on-answer alternative