973. K Closest Points to Origin

Question

LeetCode Problem

Image: View on LeetCode

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 <= 10^4
• -10^4 <= xi, yi <= 10^4

Solutions

approach 1 :
three way partition is an optimal solution for this when there are lots of same distance points.

void partition(vector<vector<int>> & points, int left, int right, int&lt, int&rt) {
	int randomPivotIdx = left + rand() % (right - left + 1);
	const auto& pv = points[randomPivotIdx];
	int px = pv[0];
	int py = pv[1];
	long long pivotDistance = 1LL * px*px + 1LL* py*py;
	lt = left;
	int mid = left;
	rt = right;
	while (mid <= rt) {
		auto const& mv = points[mid];
		int mx = mv[0];
		int my = mv[1];
		long long midDistance = 1LL * mx*mx + 1LL * my * my;
		if (midDistance < pivotDistance) {
			swap(points[mid], points[lt]);
			++mid;
			++lt;
		} else if (midDistance > pivotDistance) {
			swap(points[mid], points[rt]);
			--rt;
		} else {
			++mid;
		}
	}
}

vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
	int n = static_cast<int>(nums.size());
	if (k < 1 || k > n) return points;
	// change k to 0-based number
	int target = k - 1;
	int left = 0;
	int right = n - 1;
	while (left <= right) {
		int lt = 0;
		int rt = 0;
		partition(points, left, right, lt, rt);
		if (target < lt) {
			right = lt - 1;
		} else if (target > rt) {
			left = rt + 1;
		} else {
			return vector<vector<int>>{
				points.begin(), 
				points.begin() + k
			};
		}
	}
	// should never enter here.
	return vector<vector<int>>{};
}

approach 2:

use priority queue

   // Use a min-heap (priority_queue with custom comparator) ordered by distance
   // from the origin so that the closest points are at the top.
   // Use long long for distance calculation to avoid overflow.
   // Building the heap from the full input is O(N) and popping k elements is
   // O(k log N), so overall time complexity is O(N + k log N).
   // Space Complexity: O(N) for the heap plus O(k) for the result.
   struct Cmp {
	   bool compartor()(const vector<int>& a, const vector<int>&b) const {
		   long long distance_a = a[0] * a[0] + a[1] * a[1];
		   long long distance_b = b[0] * b[0] + b[1] * b[1];
		   return distance_a > distance_b;
	   }
   };
vector<vector<int>> kClosest(vector<vector<int>> & points, int k) {
	priority_queue<vector<int>, vector<vector<int>>, Cmp> pq(
		points.begin(), points.end()
	);
	vector<vector<int>> result;
	result.reserve(k);
	while (k-- > 0 && !pq.empty()) {
		result.push_back(pq.top());
		pq.pop();
	}
	return result;
}