LeetCode 973: K Closest Points to Origin

link

k-closest points are the k-smallest distance points. We use a max-heap to keep track of the k-smallest distance points.

Time: \mathcal{O}(n \cdot \lg{k}), space: \mathcal{O}(k).

from heapq import heappush as push, heappop as pop

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dist(a):
            x, y = a
            return x*x + y*y
        
        max_heap = []
        for p in points:
            push(max_heap, ( -dist(p), p ))
            if len(max_heap) > k:
                pop(max_heap)
        
        return [item[1] for item in max_heap]

Leave a comment