LC 973 - K Closest Points to Origin

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

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:

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:

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.

Constraints:

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

Links

Question here and solution here

Solution

concept

use a min heap to solve the problem. Take note that heap in python only see the first element if you store a tuple or list. We keep track of the distance.

code

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        min_heap = []

        def calc_dist(x, y):
            return sqrt(x**2 + y**2)

        for p in points:
            heapq.heappush(min_heap, (calc_dist(p[0], p[1]), p))

        ans = []
        while len(ans) < k and min_heap:
            dist, coord = heapq.heappop(min_heap)
            ans.append(coord)
            
        return ans

class NeetSolution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        minHeap = []
        for x, y in points:
            # if just comparison is needed, we can avoid using sqrt
            dist = (x ** 2) + (y ** 2)
            minHeap.append((dist, x, y))
        
        heapq.heapify(minHeap)
        res = []
        for _ in range(k):
            _, x, y = heapq.heappop(minHeap)
            res.append((x, y))
        return res

Complexity

time: $O(k \log n)$
space: $O(n)$