K Closest Points to Origin

Approach Summary

Max-heap of size k keyed by Euclidean distance² (no sqrt needed). Evict the farthest whenever size exceeds k.

Full Solution & Approach

The distance of a point (x, y) from the origin is the square root of x² + y², and since the square root is monotonic, comparing x² + y² is equivalent to comparing distances — so work with squared distances to stay in integers. The straightforward answer sorts all points by squared distance and takes the first k, costing O(n log n). The better answer maintains a max-heap of size k keyed by squared distance: for each point, if the heap has fewer than k entries, push it; otherwise compare its distance with the heap root — the farthest of the current k closest — and if the new point is closer, pop the root and push the new point. After the pass the heap contains exactly the k closest points. Because the heap never exceeds size k, every push and pop is O(log k), giving O(n log k) overall — a strict improvement over sorting when k is much smaller than n, and it works for streaming input. This is the same min/max-heap-of-size-k pattern used for Kth Largest Element and Top K Frequent Elements; the only twist here is keying by squared distance and using a max-heap to evict the farthest.

Each of n points is inserted into a size-k heap with O(log k) work — O(n log k) total. The heap stores at most k triples, so O(k) space.

Solution Code

Solution

import heapq

def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
    heap = []
    for x, y in points:
        d = x * x + y * y
        heapq.heappush(heap, (-d, x, y))
        if len(heap) > k:
            heapq.heappop(heap)
    return [[x, y] for _, x, y in heap]

function kClosest(points, k) {
  const heap = []; // max-heap keyed by squared distance (store negatives to keep JS sort simple)
  const push = (p) => {
    heap.push(p);
    heap.sort((a, b) => b.d - a.d);
  };
  for (const [x, y] of points) {
    const d = x * x + y * y;
    if (heap.length < k) push({ d, x, y });
    else if (d < heap[0].d) {
      heap.shift();
      push({ d, x, y });
    }
  }
  return heap.map((p) => [p.x, p.y]);
}

public int[][] kClosest(int[][] points, int k) {
    PriorityQueue<int[]> heap = new PriorityQueue<>(
        (a, b) -> (b[0] * b[0] + b[1] * b[1]) - (a[0] * a[0] + a[1] * a[1]));
    for (int[] p : points) {
        heap.offer(p);
        if (heap.size() > k) heap.poll();
    }
    return heap.toArray(new int[0][]);
}

vector<vector<int>> kClosest(const vector<vector<int>>& points, int k) {
    auto dist = [](const vector<int>& p) { return p[0] * p[0] + p[1] * p[1]; };
    priority_queue<pair<int, vector<int>>> heap; // max-heap keyed by distance
    for (const auto& p : points) {
        heap.push({dist(p), p});
        if ((int)heap.size() > k) heap.pop();
    }
    vector<vector<int>> res;
    while (!heap.empty()) { res.push_back(heap.top().second); heap.pop(); }
    return res;
}

Edge Cases to Watch

k equals the number of points — the heap retains every point
Ties in distance — any of the tied points is acceptable
Points on the axes — their squared distance is still valid
The origin itself — distance 0, always among the closest

How to Recognize This Pattern

"Find k nearest points"
Max-heap of size k beats sorting entire array

Complexity Analysis

Time Complexity

O(n log k)

Space Complexity

O(k)