leetcode-973. 最接近原点的 K 个点

题目

我们有一个由平面上的点组成的列表 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

解题思路

Heap

主要是用这道题来熟悉一下堆的应用。

首先计算出来每个点到原点的距离，然后用前K个数字建一个大小为K的大根堆，当堆建好后，后面的数字进来时，先push进堆，再pop出堆即可。

Quick sort

Find smaller pairs, larger pairs, and same pairs, then dive into only one of them.

代码

Heap

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap = []
        for x, y in points:
            if len(heap) < k:
                heapq.heappush(heap, (-x ** 2 - y ** 2, x, y))
            else:
                heapq.heappushpop(heap, (-x ** 2 - y ** 2, x, y))
        return [[item[1], item[2]] for item in heap]

Quick sort

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        dis_points = [(x ** 2 + y ** 2, x, y) for x, y in points]
        def helper(nums: list, k: int) -> list:
            if len(nums) == k:
                return nums
            pivot_index = random.randint(0, len(nums) - 1)
            pivot = nums[pivot_index][0]
            equals, smallers, largers = [], [], []
            for item in nums:
                if item[0] == pivot:
                    equals.append(item)
                elif item[0] < pivot:
                    smallers.append(item)
                else:
                    largers.append(item)
            if len(smallers) > k:
                return helper(smallers, k)
            elif len(smallers) + len(equals) < k:
                return smallers + equals + helper(largers, k - len(smallers) - len(equals))
            return smallers + equals[:k - len(smallers)]
        res = helper(dis_points, k)
        return [[item[1], item[2]] for item in res]