// input: [[3,3],[5,-1],[-2,4]] publicint[][] kClosest(int[][] points, int K) { Arrays.sort(points, (p1, p2) -> { // comparator return dis(p1) - dis(p2); // < }); int[][] result = newint[K][]; for (int i = 0; i < K; ++i) { result[i] = newint[] { points[i][0], points[i][1] }; } return result; }

Time: $O(N\log{N})$ Space: $O(K)$

Max Heap

Since we want K smallest elements, we use a K-size max heap. We repeatedly add elements to the heap. When the size is greater than K, we remove the largest element.

The idea of Quick Select is to find the K-th element, make elements on the left less than the K-th element, and make elements on the right greater than the K-th element.

We randomly pick an element. If it turns out to be the actual K-th element, we can stop at the first step. So it depends on luck.

After swapping, if this element turns out to be an element to the left of the K-th element, we recursively solve the subproblem for the right subarray; if it is to the right, we solve it for the left subarray.

Note: To make code cleaner, after we pick the random index, we place the element at the beginning before doing swapping.

// quickSelect publicint[][] kClosest(int[][] points, int K) { int n = points.length; quickSelect(points, K - 1, 0, n - 1); // index from 0 int[][] result = newint[K][]; for (int i = 0; i < K; ++i) { result[i] = points[i]; } return result; }

// find the k-th element (from 0 ~ hi - 1) privatevoidquickSelect(int[][] points, int k, int lo, int hi){ if (lo == hi) { return; } Random rand = new Random(); int randIdx = lo + rand.nextInt(hi - lo + 1); // lo + (0 ~ #element) // place the key to the beginning swap(points, lo, randIdx); int key = lo; int i = lo, j = hi + 1; // one index offset // use the quicksort template while (true) { while (dis(points[++i]) < dis(points[key])) { // move i if (i == hi) break; } while (dis(points[--j]) > dis(points[key])) { // move j if (j == lo) break; } if (i >= j) break; swap(points, i, j); } swap(points, key, j); // put [key] to the correct place [<key] [key] [>key]

// notice that k = K - 1 // j is now where [key] is if (j > k) quickSelect(points, k, lo, j - 1); // left if (j < k) quickSelect(points, k, j + 1, hi); // right // if j == k, finish. }

privatevoidswap(int[][] points, int i, int j){ int[] temp = points[i]; points[i] = points[j]; points[j] = temp; }

Time: $O(N)$ in average; $O(N^2)$ in the worst case. Space: $O(K)$