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)
// 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);
}
};
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;
}
};
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;
}
};
// 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];
}
int quickSelect(vector<vector<int>> &A, int L, int R) {
int i = L, j = L, pivotIndex = L + rand() % (R - L + 1), pivot = dist(A[pivotIndex]);
swap(A[pivotIndex], A[R]);
for (; i < R; ++i) {
if (dist(A[i]) < pivot) swap(A[i], A[j++]);
}
swap(A[j], A[R]);
return j;
}
public:
vector<vector<int>> kClosest(vector<vector<int>>& A, int K) {
int L = 0, R = A.size() - 1;
while (L < R) {
int M = quickSelect(A, L, R);
if (M + 1 == K) break;
if (M + 1 > K) R = M - 1;
else L = M + 1;
}
return vector<vector<int>>(begin(A), begin(A) + K);
}
};