973

973. K Closest Points to Origin (Medium)

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is the Euclidean distance (i.e., √(x1 - x2)2 + (y1 - y2)2).

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]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

 

Constraints:

• 1 <= k <= points.length <= 104
• -104 < xi, yi < 104

Companies:
Facebook, Amazon, DoorDash, Google, Microsoft, LinkedIn, Uber, Asana, Flipkart

Related Topics:
Array, Math, Divide and Conquer, Geometry, Sorting, Heap (Priority Queue), Quickselect

Similar Questions:

• Kth Largest Element in an Array (Medium)
• Top K Frequent Elements (Medium)
• Top K Frequent Words (Medium)
• Find Nearest Point That Has the Same X or Y Coordinate (Easy)

Solution 1. Sort

// OJ: https://leetcode.com/problems/k-closest-points-to-origin/
// Author: github.com/lzl124631x
// Time: O(NlogN)
// Space: O(1)
class Solution {
private:
    int dist(vector<int> &p) {
        return p[0] * p[0]  + p[1] * p[1];
    }
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        sort(points.begin(), points.end(), [&](auto &a, auto &b) { return dist(a) < dist(b); });
        return vector<vector<int>>(points.begin(), points.begin() + K);
    }
};

Solution 2. Heap of size N

Keep a min-heap of all elements and pop K times.

Heap creation takes O(N). Popping an element takes O(logN) which we repeat K times.

// OJ: https://leetcode.com/problems/k-closest-points-to-origin/
// Author: github.com/lzl124631x
// Time: O(N + KlogN)
// Space: O(N)
class Solution {
    int dist(vector<int> &p) {
        return p[0] * p[0] + p[1] * p[1];
    }
public:
    vector<vector<int>> kClosest(vector<vector<int>>& A, int K) {
        auto cmp = [&](auto &a, auto &b) { return dist(a) > dist(b); };
        priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> pq(begin(A), end(A), cmp);
        vector<vector<int>> ans;
        while (ans.size() < K) {
            ans.push_back(pq.top());
            pq.pop();
        }
        return ans;
    }
};

Or use make_heap and pop_heap in place.

// OJ: https://leetcode.com/problems/k-closest-points-to-origin/
// Author: github.com/lzl124631x
// Time: O(N + KlogN)
// Space: O(1)
class Solution {
    int dist(vector<int> &p) {
        return p[0] * p[0] + p[1] * p[1];
    }
public:
    vector<vector<int>> kClosest(vector<vector<int>>& A, int k) {
        auto cmp = [&](auto &a, auto &b) { return dist(a) > dist(b); };
        make_heap(begin(A), end(A), cmp);
        vector<vector<int>> ans;
        while (k--) {
            pop_heap(begin(A), end(A), cmp);
            ans.push_back(A.back());
            A.pop_back();
        }
        return ans;
    }
};

Solution 3. Heap of size K

If the space is limited, we can keep a max-heap of size K.

We loop through each point:

• If the heap has less than K elements, push the point into heap.
• Otherwise, if the distance of the element is smaller than that of the heap top, we pop the heap top and push this point into heap.

In the end, all the elements left in the heap forms the answer.

// OJ: https://leetcode.com/problems/k-closest-points-to-origin/
// Author: github.com/lzl124631x
// Time: O(NlogK)
// Space: O(K)
class Solution {
    int dist(vector<int> &p) {
        return p[0] * p[0] + p[1] * p[1];
    }
public:
    vector<vector<int>> kClosest(vector<vector<int>>& A, int K) {
        auto cmp = [&](auto &a, auto &b) { return dist(a) < dist(b); };
        priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> pq(cmp);
        for (auto &p : A) {
            pq.push(p);
            if (pq.size() > K) pq.pop();
        }
        vector<vector<int>> ans;
        while (pq.size()) {
            ans.push_back(pq.top());
            pq.pop();
        }
        return ans;
    }
};

Solution 4. Quick Select

// OJ: https://leetcode.com/problems/k-closest-points-to-origin/
// Author: github.com/lzl124631x
// Time: O(N) on average, O(N^2) in the worst case
// Space: O(1)
class Solution {
    int dist(vector<int> &p) {
        return p[0] * p[0] + p[1] * p[1];
    }
public:
    vector<vector<int>> kClosest(vector<vector<int>>& A, int k) {
        int L = 0, R = A.size() - 1;
        auto partition = [&](int L, int R) {
            int i = L, pivotIndex = L + rand() % (R - L + 1), pivot = dist(A[pivotIndex]);
            swap(A[pivotIndex], A[R]); // swap the pivot to the end of this subarray
            for (int j = L; j < R; ++j) { // traverse from L to R - 1. Note that we don't visit A[R] which is the pivot
                if (dist(A[j]) < pivot) swap(A[i++], A[j]);
            }
            swap(A[i], A[R]);
            return i;
        };
        while (true) {
            int M = partition(L, R);
            if (M + 1 == k) break;
            if (M + 1 < k) L = M + 1;
            else R = M - 1;
        }
        return vector<vector<int>>(begin(A), begin(A) + k);
    }
};