K Closest Points to Origin — LeetCode 973 Heap and Quickselect

Problem Statement#

Given an array points where points[i] = [xi, yi] and an integer k, return the k closest points to the origin (0, 0). Distance is Euclidean: sqrt(x^2 + y^2). The answer may be returned in any order.

Constraints:

• 1 <= k <= points.length <= 10^4
• -10^4 <= xi, yi <= 10^4

Input:  points = [[1, 3], [-2, 2]], k = 1
Output: [[-2, 2]]

Input:  points = [[3, 3], [5, -1], [-2, 4]], k = 2
Output: [[3, 3], [-2, 4]]

Why This Problem Matters#

LeetCode 973 K Closest Points to Origin is asked in nearly every Meta and Amazon onsite. It is a geometric reframe of "Top K by score": replace score with negative squared distance and the problem reduces to a heap or Quickselect.

The dimensional twist (2D points) hides nothing — Euclidean distance squared is monotonic with distance, so x^2 + y^2 works without the square root. Skipping sqrt is a small but impressive performance and correctness signal.

Keywords: "K closest points interview", "FAANG geometry top K", "heap Quickselect tradeoff", "L2 distance ranking".

The Core Insight#

Use squared distance to avoid floating point. Then either:

• Max-heap of size K: push and evict largest when overflow. O(n log k).
• Quickselect: partition by squared distance, recurse only into the side containing index K. O(n) average.

Senior interviewers expect both to be discussed.

Visual Dry Run#

points = [[3, 3], [5, -1], [-2, 4]], k = 2.

Result: [[3, 3], [-2, 4]].

Solution (Heap)#

import heapq
 
class Solution:
    def kClosest(self, points, k):
        h = []
        for x, y in points:
            d = -(x * x + y * y)
            if len(h) < k:
                heapq.heappush(h, (d, x, y))
            elif d > h[0][0]:
                heapq.heapreplace(h, (d, x, y))
        return [[x, y] for _, x, y in h]

class MaxHeap {
    constructor(cmp) { this.h = []; this.cmp = cmp; }
    push(v) { this.h.push(v); this._up(this.h.length - 1); }
    pop() {
        const top = this.h[0], last = this.h.pop();
        if (this.h.length) { this.h[0] = last; this._down(0); }
        return top;
    }
    peek() { return this.h[0]; }
    size() { return this.h.length; }
    _up(i) {
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.cmp(this.h[p], this.h[i]) >= 0) break;
            [this.h[p], this.h[i]] = [this.h[i], this.h[p]];
            i = p;
        }
    }
    _down(i) {
        const n = this.h.length;
        while (true) {
            const l = 2 * i + 1, r = 2 * i + 2;
            let s = i;
            if (l < n && this.cmp(this.h[l], this.h[s]) > 0) s = l;
            if (r < n && this.cmp(this.h[r], this.h[s]) > 0) s = r;
            if (s === i) break;
            [this.h[s], this.h[i]] = [this.h[i], this.h[s]];
            i = s;
        }
    }
}
 
var kClosest = function(points, k) {
    const h = new MaxHeap((a, b) => a[0] - b[0]);
    for (const [x, y] of points) {
        const d = x * x + y * y;
        if (h.size() < k) h.push([d, x, y]);
        else if (d < h.peek()[0]) {
            h.pop();
            h.push([d, x, y]);
        }
    }
    return h.h.map(([, x, y]) => [x, y]);
};

Time: O(n log k). Space: O(k).

Solution (Quickselect)#

import random
 
class Solution:
    def kClosestQS(self, points, k):
        def dist(p):
            return p[0] * p[0] + p[1] * p[1]
 
        def partition(lo, hi):
            p = random.randint(lo, hi)
            points[p], points[hi] = points[hi], points[p]
            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]

Time: O(n) average, O(n^2) worst. Space: O(1).

Common Mistakes#

• Computing sqrt — slower and risks float precision. Use squared distance.
• Using a min-heap and keeping all N points — wastes memory.
• Forgetting to handle ties — squared distance ties are fine since we want any K.
• Off-by-one in Quickselect — target index is k, not k - 1.
• Not random-pivoting Quickselect — adversarial inputs trigger O(n^2).

Interview Tips#

• Lead with the heap, then bring up Quickselect as the optimal expected-time solution.
• Mention that we square instead of sqrt — interviewers love that detail.
• Quickselect mutates the input — flag this.
• Discuss the worst-case behavior for Quickselect and how randomization fixes it.

Follow-up Questions#

• What if points stream in? Heap is the only choice — Quickselect needs the full array.
• What about K closest in higher dimensions? Same heap; replace distance with sum of squares.
• What about K closest by Manhattan distance? Replace sum-of-squares with absolute sums.
• Can we do better than O(n) average? Not without preprocessing (KD-tree, etc.).

Key Takeaways#

• LeetCode 973 reduces to "top K by smallest squared distance".
• Use squared distance to avoid sqrt and float precision.
• Max-heap of size K: O(n log k) with O(k) space, online-friendly.
• Quickselect: O(n) average, O(1) extra space, but mutates input.
• Random pivot is essential to avoid Quickselect worst case.
• Generalizes to higher dimensions and other distance metrics.
• Always discuss both approaches in a senior interview.