K Closest Points to Origin

mediumHeap / Priority Queue24 stepsJava · Python · C++

Keep the k closest points in a max-heap keyed by squared distance. When the heap grows past k, drop the farthest. What remains are the k closest.

▶ Watch the animated trace

Traced input: points = [[3, 3], [5, -1], [-2, 4], [1, 1]], k = 2 → [[3, 3], [1, 1]]. Original problem: neetcode.io/problems/k-closest-points-to-origin

Complexity

BestO(n log k)

AverageO(n log k)

WorstO(n log k)

SpaceO(k)

Compare squared distances; no square roots needed.

Key idea

To keep the k smallest, use a max-heap and evict its top.
x² + y² preserves the distance order without Math.sqrt.
Output order doesn't matter.

Common mistakes

Using sqrt and doubles: slower and introduces rounding.
Comparator overflow with a - b style comparisons on large values; use Integer.compare.

Solution

Java

int[][] kClosest(int[][] points, int k) {
    PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> Integer.compare(b[0] * b[0] + b[1] * b[1], a[0] * a[0] + a[1] * a[1]));
    for (int i = 0; i < points.length; i++) {
        heap.offer(points[i]);
        if (heap.size() > k) heap.poll();
    }
    int[][] res = new int[k][];
    for (int i = 0; i < k; i++) res[i] = heap.poll();
    return res;
}

Python

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

C++

vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
    auto dist = [](const vector<int>& p) { return p[0] * p[0] + p[1] * p[1]; };
    auto cmp = [&](const vector<int>& a, const vector<int>& b) { return dist(a) < dist(b); };
    priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> heap(cmp);
    for (auto& p : points) {
        heap.push(p);
        if ((int)heap.size() > k) heap.pop();
    }
    vector<vector<int>> res;
    while (!heap.empty()) { res.push_back(heap.top()); heap.pop(); }
    return res;
}

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
k equals n Every point is returned.	`points = [[1, 3], [-2, 2]], k = 2`	`[[1, 3], [-2, 2]]`
Origin itself Distance 0 is the closest.	`points = [[0, 0], [1, 1]], k = 1`	`[[0, 0]]`
k = 1 8 < 10.	`points = [[1, 3], [-2, 2]], k = 1`	`[[-2, 2]]`