Skip to main content
CodePeetCodePeet

Stamping The Sequence

HARDProblemSolveExternal Links

Description

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 a point and the origin is measured using the Euclidean distance formula: √(x² + y²). However, since we only need to compare distances (not compute exact values), we can compare squared distances x² + y² instead — this avoids unnecessary square root computations and floating-point precision issues.

The answer may be returned in any order. The problem guarantees that the answer is unique (except for ordering).

Examples

Example 1

Input: points = [[1, 3], [-2, 2]], k = 1

Output: [[-2, 2]]

Explanation:

  • Distance of (1, 3) from origin: √(1² + 3²) = √10 ≈ 3.16
  • Distance of (-2, 2) from origin: √((-2)² + 2²) = √8 ≈ 2.83

Since √8 < √10, the point (-2, 2) is closer to the origin. We need k = 1 closest point, so we return [[-2, 2]].

Example 2

Input: points = [[3, 3], [5, -1], [-2, 4]], k = 2

Output: [[3, 3], [-2, 4]]

Explanation:

  • Distance of (3, 3): √(9 + 9) = √18 ≈ 4.24
  • Distance of (5, -1): √(25 + 1) = √26 ≈ 5.10
  • Distance of (-2, 4): √(4 + 16) = √20 ≈ 4.47

The two closest points are (3, 3) with distance √18 and (-2, 4) with distance √20. The answer [[-2, 4], [3, 3]] would also be accepted since order does not matter.

Example 3

Input: points = [[0, 1], [1, 0]], k = 2

Output: [[0, 1], [1, 0]]

Explanation: Both points have the same distance from the origin: √(0² + 1²) = √(1² + 0²) = 1. Since k = 2 and there are exactly 2 points, we return both.

Constraints

  • 1 ≤ k ≤ points.length ≤ 10^4
  • -10^4 ≤ xi, yi ≤ 10^4
  • The answer is guaranteed to be unique (except for the order).

Editorial

Brute Force

Intuition

The most straightforward way to find the k closest points is to calculate the distance of every point from the origin, sort all the points by that distance, and then pick the first k from the sorted list.

Imagine you are a teacher and you have a map with pins representing student homes. You want to find the 3 students who live closest to the school. The simplest approach: measure the distance from the school to each pin, line them up from nearest to farthest, and pick the first 3.

Since we only need to compare relative distances (not exact values), we can compare squared distances (x² + y²) instead of actual Euclidean distances. Squaring is cheaper than computing square roots, and the ordering remains the same because the square root function is monotonically increasing.

Step-by-Step Explanation

Let's trace through with points = [[3, 3], [5, -1], [-2, 4]], k = 2:

Step 1: Compute the squared distance of each point from the origin.

  • Point [3, 3]: dist² = 3² + 3² = 9 + 9 = 18
  • Point [5, -1]: dist² = 5² + (-1)² = 25 + 1 = 26
  • Point [-2, 4]: dist² = (-2)² + 4² = 4 + 16 = 20

Step 2: Pair each point with its squared distance.

  • ([3, 3], 18), ([5, -1], 26), ([-2, 4], 20)

Step 3: Sort by squared distance in ascending order.

  • ([3, 3], 18), ([-2, 4], 20), ([5, -1], 26)

Step 4: Take the first k = 2 points.

  • Result: [[3, 3], [-2, 4]]

These are the 2 closest points to the origin.

Sort All Points by Distance — Pick First K — Watch how we compute the squared distance for each point, sort all points by distance, and select the first k elements as our answer.

Algorithm

  1. For each point [x, y], compute its squared distance from the origin: x² + y².
  2. Sort all points by their squared distance in ascending order.
  3. Return the first k points from the sorted list.

Code

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        // Sort all points by squared distance from origin
        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 the first k points
        return vector<vector<int>>(points.begin(), points.begin() + k);
    }
};
class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        # Sort all points by squared distance from origin
        points.sort(key=lambda p: p[0] ** 2 + p[1] ** 2)
        
        # Return the first k points
        return points[:k]
class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // Sort all points by squared distance from origin
        Arrays.sort(points, (a, b) -> {
            return (a[0] * a[0] + a[1] * a[1]) - (b[0] * b[0] + b[1] * b[1]);
        });
        
        // Return the first k points
        return Arrays.copyOfRange(points, 0, k);
    }
}

Complexity Analysis

Time Complexity: O(n log n)

Sorting all n points dominates the runtime. Computing the squared distance for each point is O(1), and there are n points, contributing O(n). The overall complexity is O(n log n) due to the sort.

Space Complexity: O(log n)

The sort uses O(log n) space for its internal recursion stack (Timsort in Python, introsort in C++). We do not allocate any additional data structures proportional to n beyond the output array.

Why This Approach Is Not Efficient

The sorting approach sorts all n points, even though we only need the k smallest. When k is much smaller than n (e.g., k = 10 and n = 10,000), we are doing far more work than necessary.

Sorting gives us a fully ordered list, but the problem only asks for the k smallest — it does not require them in sorted order. This is a classic case of over-solving: we computed a complete ranking when a partial one suffices.

Can we do better? Yes. Instead of sorting everything, we can use a max-heap of size k to keep track of only the k closest points seen so far. As we scan through the array, we compare each new point against the farthest point in our heap. If the new point is closer, we swap them. This way, we never sort the entire array — we only maintain a small "top-k" collection, reducing the time from O(n log n) to O(n log k).

Better Approach - Max Heap of Size K

Intuition

Instead of sorting all points, we maintain a max-heap that holds at most k points. The max-heap is ordered by distance, so the farthest of the k points is always at the top.

Think of it like a VIP lounge with exactly k seats. People (points) arrive one by one. If the lounge is not full, they walk right in. If it is full, the newcomer is compared against the person sitting farthest from the stage (origin). If the newcomer is closer, they take that seat and the farthest person leaves. At the end, the lounge contains exactly the k closest people.

Why a max-heap and not a min-heap? Because we need quick access to the farthest point among our current best k, so we can decide whether a new point should replace it. A max-heap lets us peek at (and remove) the farthest in O(log k).

This gives us O(n log k) time: we process each of the n points, and each heap operation costs O(log k). When k is much smaller than n, this is a significant improvement over O(n log n).

Step-by-Step Explanation

Let's trace through with points = [[3, 3], [5, -1], [-2, 4]], k = 2:

Step 1: Initialize an empty max-heap (ordered by squared distance, max at top).

  • Heap: empty

Step 2: Process point [3, 3]. Squared distance = 18.

  • Heap has fewer than k = 2 items. Push directly.
  • Heap: [(18, [3,3])]

Step 3: Process point [5, -1]. Squared distance = 26.

  • Heap has fewer than k = 2 items. Push directly.
  • Heap: [(26, [5,-1]), (18, [3,3])]. Max at top is 26.

Step 4: Process point [-2, 4]. Squared distance = 20.

  • Heap is full (size = k = 2). Compare new distance 20 with heap top 26.
  • 20 < 26 → the new point is closer! Replace the top.
  • Pop [5, -1] (dist² = 26), push [-2, 4] (dist² = 20).
  • Heap: [(20, [-2,4]), (18, [3,3])].

Step 5: All points processed. Extract all k points from the heap.

  • Result: [[3, 3], [-2, 4]]

Max Heap of Size K — Maintaining the K Closest Points — Watch how a max-heap of size k acts as a sliding window of the best points. New points replace the farthest in the heap only if they are closer to the origin.

Algorithm

  1. Initialize an empty max-heap.
  2. For each point [x, y] in the input array:
    a. Compute squared distance: dist = x² + y².
    b. If the heap has fewer than k elements, push (dist, point) into the heap.
    c. Otherwise, if dist is less than the top of the heap (the farthest of the current best k), pop the top and push the new point.
  3. Extract all k points from the heap and return them.

Code

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        // Max-heap: pairs of (squared_distance, point_index)
        priority_queue<pair<int, int>> maxHeap;
        
        for (int i = 0; i < points.size(); i++) {
            int dist = points[i][0] * points[i][0] + points[i][1] * points[i][1];
            
            if (maxHeap.size() < k) {
                maxHeap.push({dist, i});
            } else if (dist < maxHeap.top().first) {
                maxHeap.pop();
                maxHeap.push({dist, i});
            }
        }
        
        vector<vector<int>> result;
        while (!maxHeap.empty()) {
            result.push_back(points[maxHeap.top().second]);
            maxHeap.pop();
        }
        return result;
    }
};
import heapq

class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        # Use a max-heap of size k (negate distances for max-heap via min-heap)
        max_heap = []
        
        for x, y in points:
            dist = x * x + y * y
            
            if len(max_heap) < k:
                # Negate distance for max-heap simulation
                heapq.heappush(max_heap, (-dist, [x, y]))
            elif -dist > max_heap[0][0]:  # new point is closer
                heapq.heapreplace(max_heap, (-dist, [x, y]))
        
        return [point for _, point in max_heap]
class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // Max-heap ordered by squared distance (largest on top)
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
            (a, b) -> (b[0] * b[0] + b[1] * b[1]) - (a[0] * a[0] + a[1] * a[1])
        );
        
        for (int[] point : points) {
            maxHeap.add(point);
            if (maxHeap.size() > k) {
                maxHeap.poll(); // remove the farthest
            }
        }
        
        return maxHeap.toArray(new int[k][2]);
    }
}

Complexity Analysis

Time Complexity: O(n log k)

We iterate through all n points. For each point, we perform at most one heap insertion and one heap extraction, each costing O(log k) because the heap never exceeds size k. Total: O(n log k). When k is much smaller than n (e.g., k = 10, n = 10,000), this is significantly faster than O(n log n).

Space Complexity: O(k)

The heap stores at most k elements at any time. No other data structures proportional to n are used.

Why This Approach Is Not Efficient

The max-heap approach gives us O(n log k), which is already very good when k is small relative to n. However, there exists an even faster approach for the average case.

The heap approach still performs O(log k) work for every point in the array, even when most points are obviously not among the k closest. We are processing each point individually and maintaining a sorted structure, but the problem does not require sorted output — it only asks for any k closest points in any order.

This is actually a selection problem — we want the k-th smallest distances. The Quickselect algorithm (a variant of QuickSort) can solve selection problems in O(n) average time by partitioning the array around a pivot, similar to QuickSort, but only recursing into the partition that contains the k-th element. This avoids the overhead of maintaining a heap across all elements.

Optimal Approach - Quickselect

Intuition

Quickselect is an algorithm inspired by QuickSort that finds the k-th smallest element in an unsorted array in O(n) average time. Instead of fully sorting the array, it repeatedly partitions the array around a pivot and only recurses into the half that contains the answer.

Here is the key insight: after one partition step, the pivot is in its final sorted position. All elements to its left are smaller (or equal), and all to its right are larger. If the pivot ends up at index k, then the first k elements are the k smallest — exactly what we need, and we can stop immediately.

Think of it like organizing a group photo by height. Instead of fully sorting everyone, you pick one person as a reference (the pivot). Everyone shorter goes to the left, everyone taller goes to the right. If you need the 3 shortest people and your pivot ended up at position 3 — you are done. The 3 people on the left are the shortest, regardless of their internal order.

The algorithm works as follows:

  1. Pick a pivot element.
  2. Partition the array so elements closer than the pivot come before it, and farther elements come after.
  3. If the pivot lands at position k, we are done — the first k elements are our answer.
  4. If the pivot is at position less than k, recurse on the right half.
  5. If the pivot is at position greater than k, recurse on the left half.

On average, this runs in O(n) because each recursive call processes roughly half the remaining elements (n + n/2 + n/4 + ... ≈ 2n).

Step-by-Step Explanation

Let's trace through with points = [[3, 3], [5, -1], [-2, 4]], k = 2.
Squared distances: [3,3]→18, [5,-1]→26, [-2,4]→20.

Step 1: Set search range left = 0, right = 2. We want the k = 2 closest, so we are looking for the partition where the first 2 elements have the smallest distances.

Step 2: Choose a pivot. Let's pick the last element [-2, 4] (dist² = 20) as the pivot.

Step 3: Partition around the pivot distance 20.

  • Compare [3, 3] (dist² = 18) with pivot 20: 18 ≤ 20 → belongs in the left partition. Swap it to position 0 (stays in place). Advance store index to 1.
  • Compare [5, -1] (dist² = 26) with pivot 20: 26 > 20 → belongs in the right partition. Do not swap.

Step 4: Place the pivot at the store index (position 1). Swap [-2, 4] to index 1.

  • Array after partition: [[3,3], [-2,4], [5,-1]]
  • Pivot index = 1.

Step 5: Pivot index is 1. We need k = 2 elements (indices 0 through 1). Since pivot_index + 1 = 2 = k, the first 2 elements are exactly the 2 closest points.

Step 6: Return the first k = 2 elements: [[3, 3], [-2, 4]].

Quickselect — Partition to Find K Closest — Watch how Quickselect partitions the array around a pivot distance. Elements closer than the pivot move left, farther ones move right. One partition is enough here to isolate the k = 2 closest points.

Algorithm

  1. Define a helper function quickselect(left, right, k) that partially sorts the array so that the first k elements are the k smallest by distance.
  2. Partition step:
    a. Choose a pivot (e.g., the last element in the range).
    b. Rearrange the array so all elements with distance ≤ pivot's distance are before the pivot, and all with distance > pivot's distance are after.
    c. Let pivot_index be the pivot's final position.
  3. Recurse:
    a. If pivot_index == k - 1, we are done — the first k elements are the answer.
    b. If pivot_index < k - 1, recurse on the right half: quickselect(pivot_index + 1, right, k).
    c. If pivot_index > k - 1, recurse on the left half: quickselect(left, pivot_index - 1, k).
  4. Return the first k elements of the (partially sorted) array.

Code

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        quickselect(points, 0, points.size() - 1, k);
        return vector<vector<int>>(points.begin(), points.begin() + k);
    }
    
private:
    int dist(const vector<int>& p) {
        return p[0] * p[0] + p[1] * p[1];
    }
    
    void quickselect(vector<vector<int>>& points, int left, int right, int k) {
        if (left >= right) return;
        
        int pivotDist = dist(points[right]);
        int store = left;
        
        for (int i = left; i < right; i++) {
            if (dist(points[i]) <= pivotDist) {
                swap(points[i], points[store]);
                store++;
            }
        }
        swap(points[store], points[right]);
        
        if (store == k - 1) return;
        else if (store < k - 1) quickselect(points, store + 1, right, k);
        else quickselect(points, left, store - 1, k);
    }
};
import random

class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        def dist(point):
            return point[0] ** 2 + point[1] ** 2
        
        def quickselect(left, right):
            # Randomized pivot for better average performance
            pivot_idx = random.randint(left, right)
            points[pivot_idx], points[right] = points[right], points[pivot_idx]
            
            pivot_dist = dist(points[right])
            store = left
            
            for i in range(left, right):
                if dist(points[i]) <= pivot_dist:
                    points[i], points[store] = points[store], points[i]
                    store += 1
            
            points[store], points[right] = points[right], points[store]
            
            if store == k - 1:
                return
            elif store < k - 1:
                quickselect(store + 1, right)
            else:
                quickselect(left, store - 1)
        
        quickselect(0, len(points) - 1)
        return points[:k]
class Solution {
    public int[][] kClosest(int[][] points, int k) {
        quickselect(points, 0, points.length - 1, k);
        return Arrays.copyOfRange(points, 0, k);
    }
    
    private int dist(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }
    
    private void quickselect(int[][] points, int left, int right, int k) {
        if (left >= right) return;
        
        int pivotDist = dist(points[right]);
        int store = left;
        
        for (int i = left; i < right; i++) {
            if (dist(points[i]) <= pivotDist) {
                int[] temp = points[i];
                points[i] = points[store];
                points[store] = temp;
                store++;
            }
        }
        int[] temp = points[store];
        points[store] = points[right];
        points[right] = temp;
        
        if (store == k - 1) return;
        else if (store < k - 1) quickselect(points, store + 1, right, k);
        else quickselect(points, left, store - 1, k);
    }
}

Complexity Analysis

Time Complexity: O(n) average, O(n²) worst case

On average, Quickselect processes n + n/2 + n/4 + ... ≈ 2n elements across all recursive calls, giving O(n) average time. However, in the worst case (e.g., already sorted input with a bad pivot choice), the partition may only reduce the problem by 1 element each time, leading to O(n²). Using a randomized pivot (as shown in the Python solution) makes the worst case extremely unlikely in practice.

Space Complexity: O(log n) average

The recursion depth is O(log n) on average (each partition roughly halves the problem). In the worst case, the recursion depth could be O(n), but randomized pivoting keeps this rare. The algorithm sorts in-place, so no additional arrays are needed.