LeetCode-Python-973. 最接近原点的 K 个点（排序 + 堆 + Partition）

我们有一个由平面上的点组成的列表 points。需要从中找出 K 个距离原点 (0, 0) 最近的点。

（这里，平面上两点之间的距离是欧几里德距离。）

你可以按任何顺序返回答案。除了点坐标的顺序之外，答案确保是唯一的。

 

示例 1：

输入：points = [[1,3],[-2,2]], K = 1
输出：[[-2,2]]
解释： 
(1, 3) 和原点之间的距离为 sqrt(10)，
(-2, 2) 和原点之间的距离为 sqrt(8)，
由于 sqrt(8) < sqrt(10)，(-2, 2) 离原点更近。
我们只需要距离原点最近的 K = 1 个点，所以答案就是 [[-2,2]]。
示例 2：

输入：points = [[3,3],[5,-1],[-2,4]], K = 2
输出：[[3,3],[-2,4]]
（答案 [[-2,4],[3,3]] 也会被接受。）
 

提示：

1 <= K <= points.length <= 10000
-10000 < points[i][0] < 10000
-10000 < points[i][1] < 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/k-closest-points-to-origin
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

第一种思路：

麻瓜解，直接按距离从小到大排序，然后输出前K个点。

时间复杂度：O（NlogN）

空间复杂度：O（N）

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 x:x[0] ** 2 + x[1] ** 2)[:K]

第二种思路：

TopK 就上堆， 

可以直接建立最小堆，然后以距离为key，让所有的点入堆，

然后从最小堆里Pop K次。

本质上跟排序是一样的……

时间复杂度：O（NlogN）

空间复杂度：O（N）

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        from heapq import *
        
        distance = []
        
        heapify(distance)
        for index, (i, j) in enumerate(points):
            heappush(distance, (i ** 2 + j ** 2, index))
        res = []
        while K:
            K -= 1
            _, idx = heappop(distance)
            res.append(points[idx])
            
        return res

第三种思路：

最小堆很麻瓜，因为把所有的点都push进堆了，再pop K次取前K个最小的。

而实际上只需要 size 为 K 的堆就够了。

开一个size为K的最大堆，维护最大堆里的点为目前的前K个距离原点最近的点就够了。

遍历整个points 数组，对于每个点，

如果堆没满，新的点直接入堆，

如果堆没满，新的点和堆顶元素进行比较，

如果新的点距离更远，则说明它必定不可能是答案中包含的点。

否则，说明它有可能是答案中包含的点， 让它入堆，

然后pop掉调整后的堆顶元素，因为要维护最大堆的size为K。

时间复杂度：O（NlogK）

空间复杂度：O（K）

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        from heapq import *

        max_heap = []
        for point in points:
            x, y = point[0], point[1]

            distance = x ** 2 + y ** 2
            if len(max_heap) < K:
                heappush(max_heap, (-distance, (x, y)))
            elif -max_heap[0][0] > distance:
                heappop(max_heap)
                heappush(max_heap, (-distance, (x, y)))

        return [point for distance, point in max_heap]

第四种思路：

快排的partition的思想也可以用来解TopK问题。

时间复杂度：O（N）

空间复杂度：O（N）

class Solution(object):
    def kClosest(self, points, K):
        left, right, target = 0, len(points) - 1, K - 1
        while True:
            pos = self.partition(points, left, right)
            if pos == target:
                return points[:pos + 1]
            elif pos > K: #要往左找
                right = pos - 1
            elif pos < K: #要往右找
                left = pos + 1
                
    def partition(self, nums, left, right):
        import random
        k = random.randint(left, right)
        pivot = nums[k][0] ** 2 + nums[k][1] ** 2
        nums[left], nums[k] = nums[k], nums[left]
        index = left
        
        for i in range(left + 1, right + 1):
            if nums[i][0] ** 2 + nums[i][1] ** 2 < pivot:
                index += 1
                nums[i], nums[index] = nums[index], nums[i]
        nums[left], nums[index] = nums[index], nums[left]
        return index

Follow Up：

如果K很小，可以怎么搞？

我的想法是：

如果K非常小，比如1， 2， 3之类的，

完全可以用找数组最大数，第二大的数之类的一次扫描法来解，不需要额外开堆，

这样时间复杂度是O(N), 空间复杂度是O(1)。