so we can do a easy brute force way of sorting all the distances. But if we use heap, we can store k largest values in a max heap. We store the k smallest (closest) distances, but use a max-heap so the largest of those k is at the top (ready to be kicked out).
note that for max heap, we need to negate the value before pushing to heap.
and heapq.heappushpop() pushes the current number into heap and pops the smallest item in heap (heap[0]).
import heapq
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
heap = []
heapq.heapify(heap)
for point in points:
x,y=point
dist = -(x**2+y**2)
if len(heap)==k:
heapq.heappushpop(heap, [dist,x,y])
else:
heapq.heappush(heap,[dist,x,y])
return [i[1:] for i in heap]
n log n cuz we iter through each point and heap operations are log n. Its log k for each heap operation cuz heap is size k so is n log k.
k space