LeetCode - 973 最接近原点的 K 个点【排序】

973. 最接近原点的 K 个点

示例 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

思路

  题意：求距离原点前k近的顶点。

  方法：可直接使用快排对整个顶点数组排序，但我们只需要前k近的节点，并不需要全部节点。因此当快排的哨兵pivot刚好是第k大时即可停止快排；若pivot下标i比k小，说明已经确定好i-start+1个元素的大小但范围不够，还有k-(i-start+1)个元素未确定，则对[i+1, end]位置的元素递归排序；若pivot下标i比k大，说明确定的范围太大了，需要再对[start, i-1]排序，来确定前k个。

代码

class Solution {
    Random rand = new Random();

    public void swap(int[][] points, int i, int j) {
        int[] tmp = points[i];
        points[i] = points[j];
        points[j] = tmp;
    }

    public void partition(int[][] points, int start, int end, int k) {
        int pIndex = start + rand.nextInt(end - start + 1);
        int pivot = points[pIndex][0] * points[pIndex][0] + points[pIndex][1] * points[pIndex][1];
        swap(points, pIndex, end);
        int i = start;
        for (int j = start; j < end; ++j) {
            if (points[j][0] * points[j][0] + points[j][1] * points[j][1] < pivot) {
                swap(points, i, j);
                i++;
            }
        }
        swap(points, i, end);
        if (k < i - start + 1)
            partition(points, start, i - 1, k);
        else if (k > i - start + 1)
            partition(points, i + 1, end, k-(i-start+1));
        else 
            return;
    }

    public int[][] kClosest(int[][] points, int k) {
        partition(points, 0, points.length - 1, k);
        return Arrays.copyOfRange(points, 0, k);
    }
}

运行结果