0973.md

973. K Closest Points to Origin

Level:

Topic:

Question

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest k = 1 points from the origin, so the answer is just [[-2,2]].

Intuition

Use priority queue to keep a max heap according to its distance from the origin.

• store the index and calculate the distance between the xy and origin
• if heap size is greater than k, remove the max from heap

Code

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

public int[][] kClosest(int[][] points, int k) {
    // keep a maxheap so every poll points are the current farthest
    PriorityQueue<int[]> pq = new PriorityQueue<>((p1, p2) -> p2[1] - p1[1]);

    for (int i = 0; i < points.length; i++) {
        int dist = (int) (Math.pow(points[i][0], 2) + Math.pow(points[i][1], 2));
        pq.add(new int[]{i, dist});
        if (pq.size() > k)
            pq.poll();
    }

    int[][] ans = new int[k][2];
    for (int i = 0; i < k; i++)
        ans[i] = points[pq.poll()[0]];

    return ans;
}