K Closest Points to Origin

Description

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

Note:

1. 1 <= K <= points.length <= 10000

2. -10000 < points[i][0] < 10000

3. -10000 < points[i][1] < 10000

Solutions

Using heap

// O(k + nlog(k))
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>> &points, int K) {
        if (points.size() <= K)
            return points;
        vector<int> dist(points.size());
        for (int i = 0; i < points.size(); ++i) {
            int x = points[i][0], y = points[i][1];
            dist[i] = x * x + y * y;
        }
        auto comp = [&dist](int i, int j) {
            return dist[i] < dist[j];
        };
        vector<int> heap(K + 1);
        for (int i = 0; i < K; ++i)
            heap[i] = i;
        make_heap(heap.begin(), heap.end() - 1, comp);
        for (int i = K; i < points.size(); ++i) {
            heap[K] = i;
            push_heap(heap.begin(), heap.end(), comp);
            pop_heap(heap.begin(), heap.end(), comp);
        }
        vector<vector<int>> res;
        for (int i = 0; i < K; ++i)
            res.push_back(points[heap[i]]);
        return res;
    }
};

// O(n + klog(n))
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>> &points, int K) {
        int n = points.size();
        if (n <= K)
            return points;
        vector<int> dist(n);
        for (int i = 0; i < n; ++i) {
            int x = points[i][0], y = points[i][1];
            dist[i] = x * x + y * y;
        }
        auto comp = [&dist](int i, int j) {
            return dist[i] > dist[j];
        };
        vector<int> heap(n);
        for (int i = 0; i < n; ++i)
            heap[i] = i;
        make_heap(heap.begin(), heap.end(), comp);
        vector<vector<int>> res;
        for (int i = 0; i < K; ++i) {
            pop_heap(heap.begin(), heap.end(), comp);
            res.push_back(points[heap.back()]);
            heap.pop_back();
        }
        return res;
    }
};

Quickselect

// using STL
// O(n) on average
// O(n^2) in worst case
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>> &points, int K) {
        int n = points.size();
        if (n <= K)
            return points;
        vector<int> dist(n);
        vector<int> indices(n);
        for (int i = 0; i < n; ++i) {
            int x = points[i][0], y = points[i][1];
            dist[i] = x * x + y * y;
            indices[i] = i;
        }
        auto comp = [&dist](int i, int j) {
            return dist[i] < dist[j];
        };

        nth_element(indices.begin(), indices.begin() + K, indices.end(), comp);
        vector<vector<int>> res(K);
        for (int i = 0; i < K; ++i)
            res[i] = points[indices[i]];
        return res;
    }
};