973. K Closest Points to Origin

[#heap, #math]

## description

## 973. K Closest Points to Origin

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).

The distance between two points on the X-Y plane is the Euclidean distance (i.e., √(x1 - x2)2 + (y1 - y2)2).

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in).

Example 1:

plain text

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]].

Example 2:

plain text

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

–
1 <= k <= points.length <= 104
–
-104 <= xi, yi <= 104

Constraints:

## notes

### Intuition

Maintaining a max heap of length k will keep the k smallest elements in it.

### Implementation

Define a function to calculate the distance for a point to the origin. A small optimization we can make is using the squared distance instead of the actual Euclidean distance. Since, we calculate them all the same way and we only care about comparing their relative magnitude's against each other, this is fine. The optimization we get from this is saving on the slightly expensive square root calculation. For the distance to origin, the squared distance is just x^2 + y^2.

After this, we initialize our max heap. In Python, the heapq implementation is a min heap by default. To make it a max heap, we can simply negate the heap key (distance) for it to become a max heap.

Iterate over the points:

–
Calculate the distance for the point
–
Add the distance (negated) and the point to the heap (heapq.heappush(heap, (-d, point)))
–
If the length of the heap exceeds k, pop from the heap (this will maintain a length of k for the heap, keeping the points with the k smallest distances)

At the end, construct the list of points from the heap that have the k smallest distances and return.

–
Time: O(n log k), we iterate over all the points once and and our heap operations are worst-case, O(log k) because the heap does not exceed length k.
–
Space: O(k), only extra space created is the heap which is max length k.

### Complexity

## solution

python

from typing import List
import heapq

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        # Squared distance (for small optimization)
        def distance(point: List[int]) -> float:
            x, y = point
            return x * x + y * y
        heap = []
        for point in points:
            d = distance(point)
            heapq.heappush(heap, (-d, point))
            if len(heap) > k:
                heapq.heappop(heap)

        return [p for _, p in heap]



if __name__ == "__main__":
    # Include one-off tests here or debugging logic that can be run by running this file
    # e.g. print(solution.two_sum([1, 2, 3, 4], 3))
    solution = Solution()

--EOF--