K Closest Points to Origin

NeetCode: K Closest Points to Origin · LeetCode: 973. K Closest Points to Origin · Difficulty: Medium

Given a list of points [x, y] and an integer k, return the k points closest to the origin (0, 0) by Euclidean distance. Any order is accepted.

Example

Text

Input:  points = [[0, 2], [2, 2]], k = 1                Output: [[0, 2]]
Input:  points = [[0, 2], [2, 0], [2, 2]], k = 2        Output: [[0, 2], [2, 0]]

Comparing squared distances gives the same order and avoids square roots. A squared distance is an int value that wraps silently on overflow, so large coordinates would need long.

Brute force: sort by distance

Arrays.sort with a comparator sorts an int[][] by any key. Comparator.comparingInt is used instead of subtracting two distances, because a subtraction can overflow. clone copies the outer array, so the caller's array order is unchanged.

Java

import java.util.Arrays;
import java.util.Comparator;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        int[][] sorted = points.clone();
        Arrays.sort(sorted, Comparator.comparingInt(p -> p[0] * p[0] + p[1] * p[1]));
        return Arrays.copyOf(sorted, k);
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(Arrays.deepToString(s.kClosest(new int[][]{{0, 2}, {2, 2}}, 1)));
        System.out.println(Arrays.deepToString(s.kClosest(new int[][]{{0, 2}, {2, 0}, {2, 2}}, 2)));
    }
}

Output

[[0, 2]]
[[0, 2], [2, 0]]

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

Optimal: max-heap of size k

Keep a max-heap of the k closest points seen so far, keyed by distance. Java's PriorityQueue is a min-heap, so the comparator (a, b) -> Integer.compare(b[0], a[0]) reverses the order. Its top is the farthest of the current best k. For each new point, add it; if the heap exceeds k, poll the top, which is the farthest and cannot be in the answer. That costs O(log k) per point.

Java

import java.util.Arrays;
import java.util.PriorityQueue;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // Each entry is {squared distance, index}. The comparator puts the farthest point on top.
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> Integer.compare(b[0], a[0]));
        for (int i = 0; i < points.length; i++) {
            int[] p = points[i];
            heap.add(new int[]{p[0] * p[0] + p[1] * p[1], i});
            if (heap.size() > k) {
                heap.poll();
            }
        }
        int[][] result = new int[heap.size()][];
        int n = 0;
        for (int[] entry : heap) {
            result[n++] = points[entry[1]];
        }
        // Sorting only makes the printed order deterministic; the heap's iteration order is not sorted.
        Arrays.sort(result, (p, q) -> p[0] != q[0] ? Integer.compare(p[0], q[0]) : Integer.compare(p[1], q[1]));
        return result;
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(Arrays.deepToString(s.kClosest(new int[][]{{0, 2}, {2, 2}}, 1)));
        System.out.println(Arrays.deepToString(s.kClosest(new int[][]{{0, 2}, {2, 0}, {2, 2}}, 2)));
    }
}

Output

[[0, 2]]
[[0, 2], [2, 0]]

Time: O(n log k). Space: O(k). The final Arrays.sort only makes the printed order deterministic; the problem accepts any order. When several points tie at the k-th distance, either choice is accepted.