典型top k 問題,主要熟悉priority queue寫法。
python裡面可以用heapq。時間NlogK,空間logK。
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
pq = []
i = 0
for pt in points:
record = (-self.distsq(pt), pt)
if i < K:
heappush(pq, record)
i += 1
else:
heappushpop(pq, record)
ans = []
for i in range(len(pq)):
ans.append(heappop(pq)[1])
return ans
def distsq(self, point):
return point[0]**2 + point[1]**2
直接sort
也行,時間NlogN,常數時間,python用lambda的話只要一行
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
return sorted(points, key=lambda pt: pt[0]**2 + pt[1]**2)[0:K]