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
def kClosest(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 passes within the constraints, but it does full-sort work to answer a partial question. The interview follow-up is usually “can you do better than O(n log n)?”
Approach 2 — Max-heap of size k
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 is closer than the root, replace the root. Anything evicted can never re-enter the answer.
import heapq
from typing import List
def kClosest(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, giving O(n log k) time and O(k) space. Better than sorting whenever k ≪ n, and it also handles points arriving as a stream too large to hold in memory.
Approach 3 — Quickselect (average O(n))
We don’t need the k closest in sorted order, only 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 and farther-than-pivot, and recurses only into the side containing the k-th boundary. Each round discards a constant fraction of the array on average, giving expected linear time.
import random
from typing import List
def kClosest(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, so we stop; the first 2 slots hold distances {18, 20}, giving [[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. It uses O(1) extra space but mutates the input and doesn’t suit streaming data, which is the 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.
flowchart TD
start["Need the k closest of n"] --> stream{"Data streams or too big to hold?"}
stream -->|Yes| heap["Max-heap of size k: O(n log k)"]
stream -->|No| mutate{"Array in hand, mutation OK, any output order?"}
mutate -->|Yes| qs["Quickselect: average O(n)"]
mutate -->|No| sort["Full sort: O(n log n)"]