Solution
A straightforward solution can be to sort the given array based on a custom comparator which compares Eucledian distance of two points but that would take up O(NlogN) time.
Another approach can be using a min-heap but if you think further, it would take up O(NlogN) time as well. Instead, we will use a max heap
Approach 1: Using Max Heap
We will start storing elements in a max heap. Whenever the heap has K+1 elements, we will pop an element from the heap. The popped element is the greatest element amongst the K+1 elements so it should not be in our answer (we want the K smallest elements). At the end, the heap will contain K smallest elements. Note that to compare two elements, we have to compare them by their Eucledian distance from the origin.
vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
// by default priority queue compares vectors using the first element of the vector
priority_queue<vector<int>> pq;
for (int i=0; i<points.size(); i++) {
// store the eucledian distance as first element of the array, second element -> x coordinate, third element -> y coordinate
pq.push({points[i][0]*points[i][0] + points[i][1]*points[i][1], points[i][0], points[i][1]});
// if the heap has K+1 elements, pop an element
if (i>=K)
pq.pop();
}
vector<vector<int>> ans;
while (pq.size()>0) {
vector<int> v = pq.top();
pq.pop();
ans.push_back({v[1], v[2]});
}
return ans;
}
Consider an example:
points = [[3,3],[5,-1],[-2,4]], K = 2
i=0:
pq -> {[18,3,3]}
i=1:
pq -> {[27,5,-1], [18,3,3]}
i=2:
pq -> {[20,-2,4], [18,3,3]}
Time complexity: O(NlogK)
Space complexity: O(K)
Approach 2: Using STL algorithm nth_element
nth_element rearranges the list in such a way such that the element at the nth position is the one which should be at that position if we sort the list.
vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
nth_element(points.begin(), points.begin()+K-1, points.end(), [](vector<int> p1, vector<int> p2) {
return p1[0]*p1[0]+p1[1]*p1[1] < p2[0]*p2[0]+p2[1]*p2[1];
});
return vector<vector<int>>(points.begin(), points.begin()+K);
}
Approach 3: Using STL algorithm partial_sort
partial_sort is used for sorting not the entire range, but only a sub-part of it. Here we can partial sort K elements. Note that it considers all the elements of the array while sorting (i.e., it is different from sorting only the first K elements of array).
vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
partial_sort(points.begin(), points.begin()+K, points.end(), [](vector<int> p1, vector<int> p2) {
return p1[0]*p1[0]+p1[1]*p1[1] < p2[0]*p2[0]+p2[1]*p2[1];
});
return vector<vector<int>>(points.begin(), points.begin()+K);
}
Good watch if you’re interested in learning more about nth_element and partial_sort: CppCon 2018: Fred Tingaud “A Little Order: Delving into the STL sorting algorithms”.