Sweaty Imposter

K Closest Points to Origin

MediumTop K ElementsNeetCode 150MidOpen on LeetCodeUpdated

You're given an array points where points[i] = [x, y] is a point on a 2D plane, and an integer k. Return the k points closest to the origin (0, 0), using ordinary straight-line distance, sqrt(x^2 + y^2). You may return them in any order. The answer is guaranteed to be unique apart from its order, so no tie decides which points make the cut.

Examples

points = [[1,3],[-2,2]], k = 1 returns [[-2,2]]

(0,0)(1,3)(-2,2)
(-2, 2) is the closer of the two points to the origin.

The squared distance of (1, 3) is 1 + 9 = 10 and of (-2, 2) is 4 + 4 = 8. With k = 1, only the closer point, (-2, 2), is returned.

points = [[3,3],[5,-1],[-2,4]], k = 2 returns [[3,3],[-2,4]]

The squared distances are 18 for (3, 3), 26 for (5, -1) and 20 for (-2, 4). The two smallest are 18 and 20, and [[-2,4],[3,3]] is equally correct.

Constraints

  • 1 <= k <= points.length <= 10^4

Hints

Hint 1

Distance is the only thing you care about, could you keep the closest k seen so far without ever sorting everything?

Hint 2
A pointer toward the approach
Hint 3
The full approach, step by step

Drill This: Your Approach First

Approach

Compare squared distances x^2 + y^2, which rank points the same way as true distances without any square roots. Keep a max-heap of at most k points keyed by that value; Python's heapq is a min-heap, so the code pushes the negated distance. Whenever the heap holds more than k points, pop the farthest. The survivors are the k closest. O(n log k) time, O(k) space.

Time

O(n log k)

Space

O(k)

Pitfall

In Python, pushing the raw distance into heapq and popping when the size passes k throws away the closest point each time, leaving the k farthest. Negate the key so each pop removes the farthest point.

Solutions

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        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[][] result = new int[k][2];
        for (int i = 0; i < k; i++) result[i] = heap.poll();
        return result;
    }
}
Drill This