973. K Closest Points to Origin

Priority Queue

class Solution {
    private int getLength(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }

    public int[][] kClosest(int[][] points, int k) {
        Queue<int[]> maxHeap = new PriorityQueue<>((o1, o2) -> getLength(o2) - getLength(o1));

        for (int i = 0; i < k; i++) {
            maxHeap.offer(points[i]);
        }
        for (int i = k; i < points.length; i++) {
            if (getLength(points[i]) < getLength(maxHeap.peek())) {
                maxHeap.poll();
                maxHeap.offer(points[i]);
            }
        }

        return maxHeap.toArray(new int[0][]);
    }
}

注意需要使用大顶堆，堆底元素都比堆顶小，当后续遍历到更小元素时，弹出堆顶元素并放入堆中。

QuickSort + Two Pointers

class Solution {
    private int getLength(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }

    public int[][] kClosest(int[][] points, int k) {
        int targetIndex = k - 1;

        int left = 0, right = points.length - 1;
        while (true) {
            int index = partition(points, left, right);
            if (index < targetIndex) {
                left = index + 1;
            } else if (index > targetIndex) {
                right = index - 1;
            } else {
                return Arrays.copyOf(points, k);
            }
        }
    }

    private int partition(int[][] points, int left, int right) {
        int pivotIndex = left + ThreadLocalRandom.current().nextInt(right - left + 1);
        int pivotLength = getLength(points[pivotIndex]);
        swap(points, pivotIndex, left);

        int lt = left + 1, rt = right;
        while (true) {
            while (lt <= rt && getLength(points[lt]) < pivotLength) {
                lt++;
            }
            while (lt <= rt && getLength(points[rt]) > pivotLength) {
                rt--;
            }

            if (lt > rt) {
                break;
            }

            swap(points, lt++, rt--);
        }

        swap(points, rt, left);
        return rt;
    }

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

Reference

973. K Closest Points to Origin