K Closest Points to Origin Algorithm by using Priority Queues in C++/Java

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:
Input: points = [[1,3],[-2,2]], K = 1
Output: [[-2,2]]

Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin. We only want the closest K = 1 points from the origin, so the answer is just [[-2,2]]. Example 2: Input: points = [[3,3],[5,-1],[-2,4]], K = 2 Output: [[3,3],[-2,4]] (The answer [[-2,4],[3,3]] would also be accepted.)

Getting the K-nearest, K-shortest, K-smallest elements in an array is not difficult. You can sort the array at O(nlogn) complexity and thus return the first K elements in the sorted array. Alternatively, we can use priority queue to build a priority queue by inserting one element after another (N elements times logN complexity of rebuilding the priority queue after an element is pushed to the priority queue).

Once the priority queue is built, we then can pop out K elements in the queue, which is the answer.

K smallest elements using C++ Priority Queue

By default, the order of poping out elements from the queue (de-queue) will be from smallest to the biggest, we can write customize comparator, in C++, you can define a comparator to compare the distances to the origin (0, 0) using the following:

1
2
3

auto cmp = [](vector<int> &a, vector<int> &b) {
   return a[0]*a[0]+a[1]*a[1] > b[0]*b[0]+b[1]*b[1];   
};

auto cmp = [](vector<int> &a, vector<int> &b) {
   return a[0]*a[0]+a[1]*a[1] > b[0]*b[0]+b[1]*b[1];   
};

It is worth mentioning that the comparator looks kinda opposite (the first parameter is bigger than the second parameter), which is different than Java.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        auto cmp = [](vector<int> &a, vector<int> &b) {
            return a[0]*a[0]+a[1]*a[1] > b[0]*b[0]+b[1]*b[1];   
        };
        priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < points.size(); ++ i) {
            pq.push(points[i]);
        }
        vector<vector<int>> r;
        for (int i = 0; i < K; ++ i) {
            r.push_back(pq.top());
            pq.pop();
        }
        return r;
    }
};

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        auto cmp = [](vector<int> &a, vector<int> &b) {
            return a[0]*a[0]+a[1]*a[1] > b[0]*b[0]+b[1]*b[1];   
        };
        priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < points.size(); ++ i) {
            pq.push(points[i]);
        }
        vector<vector<int>> r;
        for (int i = 0; i < K; ++ i) {
            r.push_back(pq.top());
            pq.pop();
        }
        return r;
    }
};

K smallest elements using Java Priority Queue

In Java, the customize comparator can be defined similarly – which is a bit verbose.

1
2
3
4
5
6
7
8

Comparator<int[]> cmp = new Comparator<int[]>() {
  @Override
  public int compare(int[] a, int[] b) {
    int x = a[0]*a[0]+a[1]*a[1];
    int y = b[0]*b[0]+b[1]*b[1];
    return x < y ? -1 : ((x == y) ? 0 : 1);
  }
};

Comparator<int[]> cmp = new Comparator<int[]>() {
  @Override
  public int compare(int[] a, int[] b) {
    int x = a[0]*a[0]+a[1]*a[1];
    int y = b[0]*b[0]+b[1]*b[1];
    return x < y ? -1 : ((x == y) ? 0 : 1);
  }
};

We have to explicitly convert the boolean to integer, and the comparator defines that the first parameter is smaller than the second.

It might be possible to use the lambda expression to define the customize comparator for priority queue:

1
2
3
4
5

Comparator<int[]> cmp = (a, b) -> {
    int x = a[0]*a[0]+a[1]*a[1];
    int y = b[0]*b[0]+b[1]*b[1];
    return x < y ? -1 : ((x == y) ? 0 : 1);
};

Comparator<int[]> cmp = (a, b) -> {
    int x = a[0]*a[0]+a[1]*a[1];
    int y = b[0]*b[0]+b[1]*b[1];
    return x < y ? -1 : ((x == y) ? 0 : 1);
};

The rest are the same, pushing all elements to build the priority queue, then de-queue the K elements – which give the K-smallest elements in the array, based on the customize comparator.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        Comparator<int[]> cmp = new Comparator<int[]>() {
            @Override
            public int compare(int[] a, int[] b) {
                int x = a[0]*a[0]+a[1]*a[1];
                int y = b[0]*b[0]+b[1]*b[1];
                return x < y ? -1 : ((x == y) ? 0 : 1);
            }
        };
        PriorityQueue<int[]> pq = new PriorityQueue<>(cmp);
        for (int i = 0; i < points.length; ++ i) {
            pq.add(points[i]);
        }
        int[][] r = new int[K][];
        for (int i = 0; i < K; ++ i) {
            r[i] = pq.remove();
        }
        return r;
    }
}

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        Comparator<int[]> cmp = new Comparator<int[]>() {
            @Override
            public int compare(int[] a, int[] b) {
                int x = a[0]*a[0]+a[1]*a[1];
                int y = b[0]*b[0]+b[1]*b[1];
                return x < y ? -1 : ((x == y) ? 0 : 1);
            }
        };
        PriorityQueue<int[]> pq = new PriorityQueue<>(cmp);
        for (int i = 0; i < points.length; ++ i) {
            pq.add(points[i]);
        }
        int[][] r = new int[K][];
        for (int i = 0; i < K; ++ i) {
            r[i] = pq.remove();
        }
        return r;
    }
}

You can also use the custom sorting algorithm to find out the K closest point to the origin: K Closest Points to Origin using Custom Sorting Algorithm in C++/Java

–EOF (The Ultimate Computing & Technology Blog) —