K Closest Points to Origin

How it works

We want the k closest points, but we do NOT need them sorted.

This is a perfect use-case for QuickSelect, the same idea used in QuickSort's partition step:

• Pick a pivot point.
• Partition all points into:
  • points closer than the pivot
  • points farther than the pivot
• After partitioning, the pivot ends at its correct position in the final sorted order.
• If the pivot ends up at index p:
  • If p == k, then the left side already contains the k closest points.
  • If p < k, search the right half.
  • If p > k, search the left half.

This avoids fully sorting the array and runs in average O(N) time.

Steps

1. Define a function to compute squared distance: dist = x^2 + y^2.
2. Use a partition function:
   • Choose a pivot distance.
   • Rearrange points so all smaller distances go left, larger go right.
   • Return the pivot's final index.
3. Maintain two pointers: L = 0, R = n - 1.
4. Repeatedly partition:
   • If pivot index p == k, stop.
   • If p < k, move L = p + 1.
   • If p > k, move R = p - 1.
5. After partitioning ends, the first k points in the array are the k closest.
6. Return those k points.

Complexity

• Time complexity: O(n) in average case, O(n ^ 2) in worst case.
• Space complexity: O(1)

Things to watch for

Computing Actual Distance with Square Root

Using sqrt(x^2 + y^2) is unnecessary and introduces floating-point precision issues. Since we only compare relative distances, squared distance x^2 + y^2 preserves ordering and avoids costly square root operations.

Using the Wrong Heap Type

For the heap approach, using a min-heap requires extracting k elements at the end, while a max-heap of size k naturally keeps the k closest. Mixing these up leads to incorrect results or inefficient solutions that maintain more elements than needed.

Integer Overflow in Distance Calculation

When coordinates can be large (up to 10^4), squaring them produces values up to 10^8. While this fits in a 32-bit integer, summing two such values approaches the limit. In languages with overflow concerns, ensure your distance calculation uses appropriate integer types.

class Solution:
    def kClosest(self, points, k):
        euclidean = lambda x: x[0] ** 2 + x[1] ** 2
        def partition(l, r):
            pivotIdx = r
            pivotDist = euclidean(points[pivotIdx])
            i = l
            for j in range(l, r):
                if euclidean(points[j]) <= pivotDist:
                    points[i], points[j] = points[j], points[i]
                    i += 1
            points[i], points[r] = points[r], points[i]
            return i

        L, R = 0, len(points) - 1
        pivot = len(points)

        while pivot != k:
            pivot = partition(L, R)
            if pivot < k:
                L = pivot + 1
            else:
                R = pivot - 1
        return points[:k]

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        int L = 0, R = points.size() - 1;
        int pivot = points.size();

        while (pivot != k) {
            pivot = partition(points, L, R);
            if (pivot < k) {
                L = pivot + 1;
            } else {
                R = pivot - 1;
            }
        }
        return vector<std::vector<int>>(points.begin(), points.begin() + k);
    }

private:
    int partition(vector<vector<int>>& points, int l, int r) {
        int pivotIdx = r;
        int pivotDist = euclidean(points[pivotIdx]);
        int i = l;
        for (int j = l; j < r; j++) {
            if (euclidean(points[j]) <= pivotDist) {
                swap(points[i], points[j]);
                i++;
            }
        }
        swap(points[i], points[r]);
        return i;
    }

    int euclidean(vector<int>& point) {
        return point[0] * point[0] + point[1] * point[1];
    }
};