最接近原点的 K 个点

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

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

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

public class KClosest {
	int[][] points;

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

	private void partition(int i, int j, int K) {
		if (i >= j) {
			return;
		}
		int p1 = i;
		int p2 = j;
		int pivot = dist(ThreadLocalRandom.current().nextInt(i, j));
		while (i < j) {
			while (i < j && dist(i) < pivot) {
				i++;
			}
			while (i < j && dist(j) > pivot) {
				j--;
			}
			swap(i, j);
		}
		if (i - p1 + 1 > K) {
			partition(p1, i, K);
		} else {
			partition(i + 1, p2, K - (i - p1 + 1));
		}
	}

	public void swap(int i, int j) {
		int tmp1 = points[i][0];
		int tmp2 = points[i][1];
		points[i][0] = points[j][0];
		points[i][1] = points[j][1];
		points[j][0] = tmp1;
		points[j][1] = tmp2;
	}

	public int dist(int i) {
		return points[i][0] * points[i][0] + points[i][1] * points[i][1];
	}
}