TL;DR
Max-heap of size k on squared distance — O(n log k) time, O(k) space (quickselect gets average O(n)).
Approach 1 — Brute force (sort everything)
Compute each point’s squared distance, sort all points by it, take the first k. Squared distance preserves the ordering of Euclidean distance, so no square roots or floats are ever needed.
from typing import List
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
return sorted(points, key=lambda p: p[0] * p[0] + p[1] * p[1])[:k]
O(n log n) time, O(n) space. At n = 10^4 this comfortably passes — the constraints don’t kill it — but it does full-sort work to answer a partial question, and the follow-up (“better than O(n log n)?”) is the real interview.
Approach 2 — Max-heap of size k
The insight: maintain the k closest points seen so far. The point at risk of eviction is the farthest of the k, so keep a max-heap keyed on squared distance (negated, since heapq is a min-heap): if a new point beats the root, replace the root. Everything evicted can never be in the answer.
import heapq
from typing import List
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
heap: list[tuple[int, int, int]] = [] # (-dist², x, y)
for x, y in points:
d = x * x + y * y
if len(heap) < k:
heapq.heappush(heap, (-d, x, y))
elif -d > heap[0][0]: # closer than the current farthest
heapq.heapreplace(heap, (-d, x, y))
return [[x, y] for _, x, y in heap]
Walkthrough of points = [[3, 3], [5, -1], [-2, 4]], k = 2:
[3, 3]: d = 18, heap not full → push. Heap holds {18}.[5, -1]: d = 26, heap not full → push. Heap holds {18, 26}; root = farthest = 26.[-2, 4]: d = 20 < 26 → heapreplace evicts 26, inserts 20. Heap holds {18, 20}.- Return
[[3, 3], [-2, 4]] (order irrelevant). ✓
Complexity: each of n points costs at most one O(log k) heap operation → O(n log k) time, O(k) space. Strictly better than sorting whenever k ≪ n, and the right answer when points arrive as a stream too large to hold in memory.
Approach 3 — Quickselect (average O(n))
The insight: we don’t need the k closest in order — just partitioned to the front. Quickselect (Hoare’s selection algorithm, the select-only half of quicksort) picks a pivot, partitions the array into closer-than-pivot / farther-than-pivot, and recurses only into the side containing the k-th boundary. Each round discards a constant fraction on average, giving expected linear time.
import random
from typing import List
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
def dist(p: List[int]) -> int:
return p[0] * p[0] + p[1] * p[1]
lo, hi = 0, len(points) - 1
while lo < hi:
pivot = dist(points[random.randint(lo, hi)])
i, j, eq = lo, hi, lo
# 3-way partition: < pivot | == pivot | > pivot
while eq <= j:
d = dist(points[eq])
if d < pivot:
points[i], points[eq] = points[eq], points[i]
i += 1
eq += 1
elif d > pivot:
points[eq], points[j] = points[j], points[eq]
j -= 1
else:
eq += 1
if k <= i: # boundary inside the < block
hi = i - 1
elif k > j + 1: # boundary inside the > block
lo = j + 1
else: # k lands in the == block: done
break
return points[:k]
Walkthrough of the same example, k = 2, distances [18, 26, 20]: suppose the random pivot is [-2, 4] (pivot distance 20). The 3-way partition arranges distances as < 20 | == 20 | > 20 → [18 | 20 | 26], so i = 1, j = 1. Is k <= i? No (2 > 1). Is k > j + 1? No (2 = 2). The boundary lands in the == block → stop; the first 2 slots hold distances {18, 20} → [[3, 3], [-2, 4]]. ✓
Complexity: expected O(n) time (n + n/2 + n/4 + … with random pivots); worst case O(n²) with adversarial pivots, which randomization makes vanishingly unlikely. O(1) extra space, but it mutates the input and doesn’t suit streaming data — the classic trade-off against Approach 2.
Common pitfalls
- Computing
math.sqrt — needless float work, and float rounding can (in other problems) corrupt comparisons; squared distance is a monotone stand-in.
- Building a min-heap of size k instead of a max-heap: a min-heap’s root is the closest point, so you can’t tell whether a new point should evict anything. For “k smallest keep-set,” the heap must expose the largest member.
- Quickselect without randomized (or median-of-medians) pivots — sorted or duplicate-heavy input degrades to O(n²); the 3-way partition specifically handles duplicate distances.
- Off-by-one on the boundary:
k counts points, partition indices count positions — after partitioning, the answer is points[:k] only once the k-th boundary falls inside the equal-to-pivot block or the loop narrows to it.
Pattern takeaway
“k smallest/closest/cheapest out of n” has a standard ladder: full sort O(n log n) → bounded heap O(n log k) → quickselect average O(n). Choose the heap when k ≪ n or the data streams; choose quickselect when the whole array is in hand, mutation is allowed, and order within the answer doesn’t matter. Either way, compare with squared (or otherwise monotone-transformed) keys instead of computing expensive exact ones.