The problem
Given a list of points [x, y] on a plane and a number k, return the k points closest to the origin (0, 0), in any order. The answer is guaranteed to be unique.
Examples
01
- Input
points = [[1, 3], [-2, 2]], k = 1
- Output
[[-2, 2]]
02
- Input
points = [[3, 3], [5, -1], [-2, 4]], k = 2
- Output
[[3, 3], [-2, 4]]
Constraints
- 1 ≤ k ≤ points.length ≤ 10⁴
- −10⁴ ≤ x, y ≤ 10⁴
The idea
The distance from the origin is √(x² + y²). Taking a square root never changes which of two numbers is bigger, so compare x² + y² instead: for (1, 3) that is 1 + 9 = 10, for (−2, 2) it is 4 + 4 = 8, so (−2, 2) is closer.
Keep the k closest points so far in a max-heap by that squared distance: the farthest of them is on top. Add each point; when there are k + 1, remove the top — the farthest cannot be among the k closest. What remains is the answer.
- Time
- O(n log k)
- Space
- O(k)
Solution · every language run against every case
class Solution: def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]: # A max-heap of the k closest so far (distances negated for Python's min-heap). # Squared distance x² + y² orders points the same as distance, without a square root. heap = [] for x, y in points: heappush(heap, (-(x * x + y * y), x, y)) if len(heap) > k: heappop(heap) # the farthest of k + 1 is not among the closest k return [[x, y] for _, x, y in heap]