演算法練習 Leetcode 力扣 973. K Closest Points to Origin解法

典型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]
        

Leave a Comment

Your email address will not be published.