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)$