973. K Closest Points to Origin
1class Solution:
2 def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
3 # heap will work
4 heap = []
5 for i, coord in enumerate(points):
6 x, y = coord
7 heapq.heappush(heap, ((x**2 + y**2) ** 0.5, i)) # (dist, idx in points)
8
9 result = []
10 for i in range(k):
11 dist, i_points = heapq.heappop(heap)
12 result.append(points[i_points])
13
14 return result
we can also do a max-heap with capacity of k (we need max heap so that we can pop the furthest one without going through all the points).
1class Solution:
2 def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
3 # heap will work
4 heap = []
5 for x, y in points:
6 heapq.heappush(heap, (-(x * x + y * y), x, y)) # (dist**2, x, y)
7 if len(heap) > k:
8 heapq.heappop(heap)
9
10 return [[x, y] for _, x, y in heap]