Reference: LeetCode EPI 10.4
Difficulty: Medium

My Post: Java Solutions with Exp. & Comments (Sorting, Heap, QuickSelect)

Problem

We have a list of points on the plane. Find the K closest points to the origin (0, 0).

(Here, the distance between two points on a plane is the Euclidean distance.)

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in.)

Example:

1
2
Input: points = [[1,3],[-2,2]], K = 1
Output: [[-2,2]]
1
2
Input: points = [[3,3],[5,-1],[-2,4]], K = 2
Output: [[3,3],[-2,4]]

Analysis

Sorting

First let’s define the distance calculation method.

1
2
3
private int dis(int[] p) { // square
return p[0] * p[0] + p[1] * p[1];
}
1
2
3
4
5
6
7
8
9
10
11
// input: [[3,3],[5,-1],[-2,4]]
public int[][] kClosest(int[][] points, int K) {
Arrays.sort(points, (p1, p2) -> { // comparator
return dis(p1) - dis(p2); // <
});
int[][] result = new int[K][];
for (int i = 0; i < K; ++i) {
result[i] = new int[] { 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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
public int[][] kClosest(int[][] points, int K) {
if (K == 0 || points == null || points.length == 0) {
return new int[0][];
}
PriorityQueue<Integer> maxPQ = new PriorityQueue<>((o1, o2) -> {
int[] p1 = points[o1];
int[] p2 = points[o2];
return dis(p2) - dis(p1);
});
int n = points.length;
for (int i = 0; i < n; ++i) {
maxPQ.offer(i);
if (maxPQ.size() > K) {
maxPQ.poll();
}
}
// output
int[][] output = new int[K][];
for (int i = K - 1; i >= 0; --i) {
output[i] = points[maxPQ.poll()];
}
return output;
}

private int dis(int[] p) {
return p[0] * p[0] + p[1] * p[1];
}

We can also write the comparator like this.

1
2
3
4
Comparator<int[]> comp = (p1, p2) -> (
return dis(p2) - dis(p1);
);
PriorityQueue<int[]> pq = new PriorityQueue<>(comp);

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

Quick Select

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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
// quickSelect
public int[][] kClosest(int[][] points, int K) {
int n = points.length;
quickSelect(points, K - 1, 0, n - 1); // index from 0
int[][] result = new int[K][];
for (int i = 0; i < K; ++i) {
result[i] = points[i];
}
return result;
}

// find the k-th element (from 0 ~ hi - 1)
private void quickSelect(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.
}

private void swap(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)$


Comment