318. K Closest Points to Origin

Keep a max-heap holding the k closest points seen so far, keyed by squared distance. When it exceeds k, pop the farthest. No need for actual square roots — comparing squared distances preserves order.

Squared distance is enough

Euclidean distance is monotonic in the squared distance x² + y², so we compare x² + y² and skip the costly sqrt.

Max-heap of size k

Maintain a heap of at most k points, ordered so the farthest is on top. For each point: push it; if the heap grows beyond k, pop the farthest. After processing all points the heap holds exactly the k closest.

for each point:
    push(-dist, point)   # negate for a max-heap on a min-heap library
    if size > k: pop()
return everything in the heap

Worked example

points = [[3,3],[5,-1],[-2,4]], k = 2: squared distances 18, 26, 20 → the two closest are (3,3) and (-2,4).

Complexity

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

import java.util.PriorityQueue;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // Max-heap by squared distance: the farthest point sits on top.
        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();
        }
        int[][] res = new int[k][2];
        for (int i = 0; i < k; i++) {
            res[i] = heap.poll();
        }
        return res;
    }
}