K Closest Points to Origin

1public int[][] kClosest(int[][] points, int k) {
2    // Max-heap to store k closest points
3    // Priority queue with custom comparator
4    PriorityQueue<int[]> heap = new PriorityQueue<>(
5        (a, b) -> Double.compare(b[2], a[2])
6    );
7
8    for (int i = 0; i < points.length; i++) {
9        int x = points[i][0];
10        int y = points[i][1];
11        double distance = Math.sqrt(x * x + y * y);
12
13        if (heap.size() < k) {
14            // Heap not full, add point
15            heap.offer(new int[]{x, y, (int)(distance * 1000)});
16        } else {
17            // Heap full, compare with max
18            double maxDistance = heap.peek()[2] / 1000.0;
19
20            if (distance < maxDistance) {
21                // Current point is closer, replace max
22                heap.poll();
23                heap.offer(new int[]{x, y, (int)(distance * 1000)});
24            }
25        }
26    }
27
28    // Extract result
29    int[][] result = new int[k][2];
30    int idx = 0;
31    while (!heap.isEmpty()) {
32        int[] point = heap.poll();
33        result[idx++] = new int[]{point[0], point[1]};
34    }
35    return result;
36}