K Closest Points to Origin
MedPrompt
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
Approaches
3 approaches with complexity analysis and trade-offs.
The most straightforward approach is to calculate the distance of each point from the origin, sort all the points based on this distance, and then pick the first k points from the sorted list.
Algorithm
- Calculate the squared Euclidean distance (
x*x + y*y) for each point to avoid using square roots. - Sort the entire
pointsarray using a custom comparator that compares points based on their calculated squared distances. - Create a new result array and copy the first
kelements from the sorted array into it.
Walkthrough
We can calculate the Euclidean distance for each point from the origin (0,0). A key optimization is to use the squared Euclidean distance (x^2 + y^2) instead of the actual distance (sqrt(x^2 + y^2)). This avoids floating-point arithmetic and the costly square root operation, while preserving the order of distances. After calculating the squared distance for every point, we can sort the entire array of points using a custom comparator that compares these squared distances. Once the array is sorted in ascending order of distance, the first k elements are our answer.
import java.util.Arrays; class Solution { public int[][] kClosest(int[][] points, int k) { // Sort the array with a custom lambda comparator. // The comparator calculates the squared distance from the origin for two points // and compares them. Arrays.sort(points, (p1, p2) -> (p1[0] * p1[0] + p1[1] * p1[1]) - (p2[0] * p2[0] + p2[1] * p2[1])); // Return the first k elements of the sorted array. // Arrays.copyOfRange is a convenient way to do this. return Arrays.copyOfRange(points, 0, k); }}Complexity
Time
O(N log N), where N is the number of points. The dominant operation is sorting the entire array.
Space
O(log N) to O(N). The space complexity depends on the implementation of the sorting algorithm. In Java, `Arrays.sort` for object arrays uses Timsort, which has a worst-case space complexity of O(N). For primitive arrays, it uses a variant of quicksort with O(log N) space on average.
Trade-offs
Pros
Simple to understand and implement.
Leverages built-in, highly optimized sorting functions from the standard library.
Cons
This approach is not the most efficient because it does more work than necessary by sorting the entire array of
Npoints, even when we only need theksmallest elements.The performance degrades as the number of points
Nincreases, regardless of how smallkis.
A more optimized approach uses a max heap (or a max priority queue) to keep track of the k closest points found so far. By maintaining a heap of size k, we can process each point in O(log k) time, which is more efficient than a full sort when k is much smaller than N.
Algorithm
- Initialize a max priority queue (max heap) which will store points, ordered by their squared distance from the origin in descending order.
- Iterate through each point in the input array.
- For each point, add it to the max heap.
- If the heap's size becomes greater than
k, remove the top element (which is the point with the largest distance currently in the heap). - After iterating through all points, the heap will contain the
kclosest points. - Extract these
kpoints from the heap and return them as an array.
Walkthrough
We iterate through the points one by one. We use a max heap to store the k points that are closest to the origin among the points we have visited so far. The heap is ordered by the squared distance from the origin, with the point farthest from the origin at the top (the 'max' element).
For each point, we compare its distance with the distance of the point at the top of the heap. If the heap isn't full (has fewer than k elements), we simply add the current point. If the heap is full and the current point is closer to the origin than the point at the top of the heap, we remove the top element (the farthest one) and insert the current point. This ensures the heap always contains the k closest points seen up to that moment.
After iterating through all N points, the heap will hold the k closest points in the entire dataset.
import java.util.PriorityQueue; class Solution { public int[][] kClosest(int[][] points, int k) { // Create a max heap (PriorityQueue in Java). // The comparator is set up to create a max heap based on squared distance. // (p2_dist - p1_dist) for max heap. PriorityQueue<int[]> maxHeap = new PriorityQueue<>((p1, p2) -> (p2[0] * p2[0] + p2[1] * p2[1]) - (p1[0] * p1[0] + p1[1] * p1[1]) ); for (int[] point : points) { maxHeap.offer(point); if (maxHeap.size() > k) { maxHeap.poll(); // Remove the farthest point. } } // The heap now contains the k closest points. // Convert the heap to a 2D array. int[][] result = new int[k][2]; for (int i = 0; i < k; i++) { result[i] = maxHeap.poll(); } return result; }}Complexity
Time
O(N log k), where N is the total number of points and k is the number of closest points to find. For each of the N points, we perform an offer and potentially a poll operation on the heap, which takes O(log k) time.
Space
O(k), as we need to store at most `k` points in the priority queue.
Trade-offs
Pros
More efficient than sorting the entire array, especially when
kis much smaller thanN.Processes points in a single pass without needing to hold all
Npoints in a sorted structure.
Cons
Can be slower than the sorting approach if
kis very close toN, as thelog kfactor approacheslog N.Slightly more complex to set up the data structure compared to a simple sort.
The most optimal approach on average is to use a selection algorithm like Quickselect. This algorithm finds the k-th smallest element in an unordered list in average linear time, effectively partitioning the array into two parts: the k closest points and the rest, without fully sorting the array.
Algorithm
- Implement a partitioning scheme similar to Quicksort. The partition function will rearrange a subarray around a pivot based on distance from the origin.
- Start with the entire array (
left = 0,right = N-1). - In a loop, partition the current range
[left, right]and get the pivot's final index,pivotIndex. - If
pivotIndexis equal tok-1, we have found the k-th smallest element. The elements from index 0 tok-1are thekclosest points. - If
pivotIndexis greater thank-1, the k-th element must be in the left subarray, so updateright = pivotIndex - 1. - If
pivotIndexis less thank-1, the k-th element must be in the right subarray, so updateleft = pivotIndex + 1. - Repeat until
pivotIndex == k-1. - Return the first
kelements of the now partially sorted array.
Walkthrough
Quickselect is a selection algorithm that modifies the Quicksort algorithm to find the k-th smallest element. The core idea is to repeatedly partition the input array. In each step, we choose a pivot element and partition the array around it, such that all elements smaller than the pivot are on its left and all elements larger are on its right. After partitioning, the pivot is in its final sorted position, say at index p.
We then compare p with k.
- If
p == k-1, we have found the k-th element, and the firstkelements of the array are thekclosest points. - If
p > k-1, the k-th closest point must be in the left subarray, so we recurse on the left part. - If
p < k-1, it must be in the right subarray, so we recurse on the right.
This process avoids fully sorting the entire array, leading to a better average time complexity. The comparison is done using the squared Euclidean distance to avoid sqrt.
import java.util.Arrays; class Solution { public int[][] kClosest(int[][] points, int k) { int left = 0, right = points.length - 1; // We are looking for the k-th element, which will be at index k-1 int targetIndex = k - 1; while (left <= right) { int pivotIndex = partition(points, left, right); if (pivotIndex == targetIndex) { break; } else if (pivotIndex > targetIndex) { right = pivotIndex - 1; } else { left = pivotIndex + 1; } } return Arrays.copyOfRange(points, 0, k); } private int partition(int[][] points, int left, int right) { // Using the rightmost element as the pivot int[] pivot = points[right]; int pivotDist = dist(pivot); int i = left; for (int j = left; j < right; j++) { if (dist(points[j]) <= pivotDist) { swap(points, i, j); i++; } } swap(points, i, right); return i; } private int dist(int[] point) { return point[0] * point[0] + point[1] * point[1]; } private void swap(int[][] points, int i, int j) { int[] temp = points[i]; points[i] = points[j]; points[j] = temp; }}Complexity
Time
Average Case: O(N). The algorithm, on average, discards half of the array in each partitioning step. Worst Case: O(N^2). This occurs with a poor choice of pivots, leading to highly unbalanced partitions.
Space
O(1) for the iterative implementation. A recursive implementation would have an average space complexity of O(log N) and a worst-case of O(N) for the recursion stack.
Trade-offs
Pros
Optimal average time complexity of O(N).
Very efficient in terms of space, as it can be implemented iteratively to use O(1) space.
Cons
Has a worst-case time complexity of O(N^2), which can occur with consistently poor pivot choices, although this is rare in practice.
The implementation is more complex than the sorting or heap-based approaches.
This approach modifies the input array in-place.
Solutions
Solution
class Solution {public int[][] kClosest(int[][] points, int k) { Arrays.sort( points, (a, b)->{ int d1 = a[0] * a[0] + a[1] * a[1]; int d2 = b[0] * b[0] + b[1] * b[1]; return d1 - d2; }); return Arrays.copyOfRange(points, 0, k); }}class Solution {public: vector<vector<int>> kClosest(vector<vector<int>> &points, int k) { sort(points.begin(), points.end(), [](const vector<int> &a, const vector<int> &b) { return a[0] * a[0] + a[1] * a[1] < b[0] * b[0] + b[1] * b[1]; }); return vector<vector<int>>(points.begin(), points.begin() + k); }};class Solution: def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]: points . sort(key=lambda p: p[0] * p[0] + p[1] * p[1]) return points[: k]Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.