K Closest Points to Origin

Ссылка на leetcode: https://leetcode.com/problems/k-closest-points-to-origin/description/

Код решения при помощи QuickSelect:

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        int start = 0;
        int end = points.length - 1;
        while (true) {
            if (start < end) {
                swap(points, start, ThreadLocalRandom.current().nextInt(start, end));
            }
            int left = start + 1;
            int right = end;
            int pivotIndex = start;
            int pivot[] = points[pivotIndex];
            double pivotDistance = distance(pivot);
            while (left <= right) {
                if (distance(points[left]) > pivotDistance && distance(points[right]) < pivotDistance) {
                    swap(points, left, right);
                    left++;
                    right--;
                }
                if (distance(points[left]) <= pivotDistance) left++;
                if (distance(points[right]) >= pivotDistance) right--;
            }
            swap(points, pivotIndex, right);
            if (right == k - 1) break;
            if (right > k - 1) {
                end = right - 1;
            } else {
                start = right + 1;
            }
        }
        int[][] result = new int[k][2];
        for (int i = 0; i < k; i++) {
            result[i][0] = points[i][0];
            result[i][1] = points[i][1];
        }
        return result;
    }

    private double distance(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }

    private void swap(int[][] points, int i, int j) {
        int tmpX = points[i][0];
        int tmpY = points[i][1];
        points[i][0] = points[j][0];
        points[i][1] = points[j][1];
        points[j][0] = tmpX;
        points[j][1] = tmpY;
    } 
}