M.A.T.R.I.X

Appearance

4.17.3 — K Closest Points to Origin ​

LeetCode 973 · Medium

The problem ​

Return the k points closest to the origin. Any order is fine.

points = [[1,3],[-2,2]], k = 1   →  [[-2,2]]

Up to 10,000 points.

The pattern ​

Two decisions, and the first one is free.

Do not take the square root. Distance is \sqrt{x^2 + y^2}, but you are only comparing distances, and the square root is monotonic — it preserves order. So compare x^2 + y^2 directly. That avoids a floating-point operation per point and, more importantly, avoids floating-point rounding making two genuinely different distances compare equal.

Comparing squared distances instead of distances is a habit worth having, and it appears anywhere geometry meets sorting.

Then: a size-k heap. You want the k smallest distances, so the item you evict is the largest of the ones you are keeping — which means a max-heap.

Note the flip from 4.17.1, where you wanted the k largest and used a min-heap. The rule is the same in both:

The root must be the item you are about to throw away.

k largest → min-heap. k smallest → max-heap.

The solution ​

python
import heapq

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap = []                                        # max-heap via negation

        for x, y in points:
            d = x * x + y * y                            # no square root
            heapq.heappush(heap, (-d, x, y))
            if len(heap) > k:
                heapq.heappop(heap)                      # drop the farthest

        return [[x, y] for _, x, y in heap]
ts
function kClosest(points: number[][], k: number): number[][] {
  // MaxHeap keyed by squared distance
  const heap = new MaxHeap<number[]>((p) => p[0] * p[0] + p[1] * p[1]);
  for (const p of points) {
    heap.push(p);
    if (heap.size > k) heap.pop();
  }
  return heap.toArray();
}

Negating the distance turns Python's min-heap into a max-heap, so the root is the farthest of the k kept so far — exactly the one a closer point should replace.

Complexity ​

O(n \log k) time, O(k) space.

The three solutions, and how to choose ​

Sort by distance and take the first k — O(n \log n), three lines, and completely acceptable:

python
return sorted(points, key=lambda p: p[0]**2 + p[1]**2)[:k]

Say this first. It is correct and it establishes the squared-distance point immediately.

Size-k heap — O(n \log k), and it works on a stream, since it never needs all the points at once.

Quickselect — O(n) on average. Partition around a pivot as in quicksort, but recurse into only the side containing the k-th position. After partitioning, the first k elements are the k closest, in no particular order — which is all the problem asks for.

Quickselect's worst case is O(n^2) if pivots are chosen badly, fixed in practice by choosing a random pivot. It also reorders the input.

timespaceworks on a stream
sortO(n \log n)O(n)no
size-k heapO(n \log k)O(k)yes
quickselectO(n) averageO(1)no

The interview answer is to name all three and the condition that picks each. If they push for the best asymptotic time, quickselect. If they mention a stream or limited memory, the heap.

Quickselect, briefly ​

python
def kClosest(self, points, k):
    def dist(p): return p[0]**2 + p[1]**2

    def partition(lo, hi) -> int:
        pivot = dist(points[hi])
        store = lo
        for i in range(lo, hi):
            if dist(points[i]) < pivot:
                points[store], points[i] = points[i], points[store]
                store += 1
        points[store], points[hi] = points[hi], points[store]
        return store

    lo, hi = 0, len(points) - 1
    while lo < hi:
        p = partition(lo, hi)
        if p == k: break
        if p < k: lo = p + 1
        else: hi = p - 1

    return points[:k]

Why it is O(n) on average: the first partition costs n, the second costs about n/2, then n/4, and so on. That sum is bounded by 2n. Contrast quicksort, which recurses into both halves and pays O(n) per level across \log n levels.

Randomising the pivot before partitioning is what keeps the average case honest on sorted or adversarial input.

Where this goes next ​

  • Kth Largest Element in an Array — the same three solutions on plain numbers. 4.17.4.
  • Top K Frequent Elements — same again, over counts, with a bucket-sort option because frequencies are bounded. 4.4.5.
  • Nearest neighbour search at scale — a k-d tree or a locality-sensitive hash, because a linear scan over billions of points is too slow no matter which selection method follows it.

What the interviewer will push on ​

"Why no square root?" It is monotonic, so it does not change the ordering, and skipping it avoids floating-point error and work.

"Max-heap or min-heap?" Max-heap for the k smallest — the root is what you evict.

"Can you do better than O(n \log k)?" Quickselect, O(n) average, and say why: the partition sizes halve, summing to 2n.

"What if the points arrive as a stream?" Only the heap survives; quickselect needs the whole array.

"What if k is 1?" A single linear scan for the minimum. Worth noticing that the general machinery collapses.

One thing to volunteer: state the squared-distance point before writing anything. It is a small thing that immediately signals you have thought about the geometry rather than copied a template.

Next: 4.17.4 Kth Largest Element in an Array — the same selection question, stripped to its simplest form.