AlgoDesign

K closest stars

1import math
2from heapq import *
3class Solution:
4    def kClosestStars(self, stars, k):
5        maxHeap = []
6        for i in range(k): 
7            distance = math.sqrt((stars[i][0])**2 + (stars[i][1])**2 + (stars[i][2])**2)
8            heappush(maxHeap, (-distance, stars[i]))
9        
10        for i in range(k, len(stars)):
11            distance = math.sqrt((stars[i][0])**2 + (stars[i][1])**2 + (stars[i][2])**2)
12            if distance < -maxHeap[0][0]:
13                heappop(maxHeap)
14                heappush(maxHeap, (-distance, stars[i]))
15        
16        res = []
17        while maxHeap:
18            _, star = heappop(maxHeap)
19            res.append(star)
20        return res