Start learning free

LESSON 392 OF 543 · HEAP

K Closest Points to Origin in Rust: Max-Heap of Size k, O(n log k)

Practice O(n log k) Read the full explanation

Given an array of points, where points[i] = [x, y] is a point on the plane, and an integer k, return the k points closest to the origin (0, 0).

The distance between a point [x, y] and the origin is the usual straight-line distance, √(x2 + y2).

Return the points sorted by their distance to the origin, the closest first. Points at the same distance come in order of their x, and points with the same x too in order of their y. The answer is guaranteed to be unique: every point left out is farther from the origin than every point in the answer.

 

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation: [1,3] is √10 away from the origin and [-2,2] is √8 away.
√8 is smaller, so [-2,2] is the closest point.

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The squared distances are 18, 26 and 20. The two closest points are [3,3] (18) and [-2,4] (20), in that order.

Example 3:

Input: points = [[0,2],[2,0],[1,1],[0,-2],[3,3]], k = 4
Output: [[1,1],[0,-2],[0,2],[2,0]]
Explanation: [1,1] is the closest (√2). [0,2], [2,0] and [0,-2] are all 2 away, so they come by x:
[0,-2] and [0,2] (x = 0) before [2,0], and [0,-2] before [0,2] because -2 < 2. [3,3] is the farthest and is left out.

 

Constraints:

  • 1 <= k <= points.length <= 105
  • -104 <= x, y <= 104
  • The same point can appear more than once.
  • The answer is unique.

Note:

Your algorithm should run in O(n log k) time and use O(k) extra space, where n is the number of points. Sorting all the points by distance would take O(n log n).

No square roots are needed: one point is closer than another exactly when its x2 + y2 is smaller, and that squared distance is a whole number of at most 2 · 108, which fits in a 32-bit integer.

STEP 1 OF 8
    Ctrl + Enter to check

    RUST · HEAP

    K Closest Points To Origin in Rust, explained

    Problem overview

    k_closest receives points: Vec<Vec<i32>>, where each point is [x, y], and an i32 named k. It returns the k points closest to the origin (0, 0), sorted by distance with the closest first. Points at the same distance are ordered by x, and then by y. The answer is unique: every point left out is farther away than every point in it.

    For points = [[3, 3], [5, -1], [-2, 4]] and k = 2, the squared distances are 18, 26 and 20, so the answer is [[3, 3], [-2, 4]]. For [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]] and k = 4 the answer is [[1, 1], [0, -2], [0, 2], [2, 0]]: [1, 1] is closest, the three points at squared distance 4 follow ordered by x and then y, and [3, 3] is left out.

    Approach: keep only the k best

    You could sort all the points by distance and take the first k, which costs O(n log n). But you only need the k closest, so you do not need the rest in order.

    Two ideas make the work smaller:

    • Compare squared distances. A point is closer than another exactly when its x * x + y * y is smaller, so no square root is needed. The value stays a whole number, and with coordinates up to 10,000 in size it is at most 200,000,000, which fits in an i32.
    • A max-heap of size k. Keep the best k points seen so far in a max-heap, where the biggest entry is on top. That top entry is the farthest point you are keeping, so it is the one to drop when a closer point arrives. Push every point, and as soon as the heap holds k + 1 entries, pop the top. Among k + 1 points, the farthest cannot be among the k closest.

    Each push and pop costs O(log k), and there are n of them, so the total is O(n log k).

    Step by step

    1. Declare the heap: let mut heap: BinaryHeap<(i32, i32, i32)> = BinaryHeap::new();. Each entry is the tuple (dist, x, y). Tuples compare item by item, so entries are ordered by distance first, then x, then y. BinaryHeap is a max-heap, so the largest tuple is on top.
    2. Loop over &points. The & borrows the vector, so the loop reads the points without taking them over.
    3. Name the coordinates: let (x, y) = (point[0], point[1]);.
    4. Compute the squared distance: let dist = x * x + y * y;.
    5. Push (dist, x, y) onto the heap. The entry keeps the point together with the number the heap orders by.
    6. Trim the heap: if heap.len() > k as usize, call heap.pop(). len() is a usize, so k is converted first. After each point, the heap holds the k closest points seen so far.
    7. Sort the survivors: heap.into_sorted_vec() returns the entries in a vector, from smallest to biggest. Because tuples sort by distance, then x, then y, this is exactly the order the answer needs.
    8. Drop the distances: closest.into_iter().map(|(_, x, y)| vec![x, y]).collect() turns each tuple into a two-number vector. The _ ignores the distance. This is the last expression, with no semicolon, so it is the return value.

    For the second example, the heap goes like this. After the first two points it holds (18, 3, 3) and (26, 5, -1), with (26, 5, -1) on top. The third point, (20, -2, 4), makes the size 3, so the top, (26, 5, -1), is popped. The heap ends with (18, 3, 3) and (20, -2, 4), and sorting them gives [[3, 3], [-2, 4]].

    The solution in Rust

    use std::collections::BinaryHeap;

impl Solution {
    fn k_closest(points: Vec<Vec<i32>>, k: i32) -> Vec<Vec<i32>> {
        let mut heap: BinaryHeap<(i32, i32, i32)> = BinaryHeap::new();
        for point in &points {
            let (x, y) = (point[0], point[1]);
            let dist = x * x + y * y;
            heap.push((dist, x, y));
            if heap.len() > k as usize {
                heap.pop();
            }
        }
        let closest = heap.into_sorted_vec();
        closest.into_iter().map(|(_, x, y)| vec![x, y]).collect()
    }
}

    The Rust details worth a closer look:

    • Max-heap by default. BinaryHeap pops the largest element. That is what this problem wants: the largest entry is the farthest point, the one to throw away. For a min-heap you would wrap values in std::cmp::Reverse, which is not needed here.
    • Tuples are ordered. A tuple of integers implements Ord, comparing the first items, then the second on a tie, and so on. That one fact gives both the heap's order and the final sort order, including the tie-breaks by x and y.
    • Type annotation. BinaryHeap<(i32, i32, i32)> states the entry type. The import use std::collections::BinaryHeap; is already in the lesson's starter code.
    • into_sorted_vec(). It consumes the heap and returns a Vec in ascending order. It does not return pops in order: the smallest comes first.
    • Casting k. k as usize is fine because k is at least 1. A negative i32 would turn into a huge usize.
    • Closure pattern. |(_, x, y)| takes a tuple apart in the closure's parameter list. into_iter() moves the tuples out of the vector, so the closure receives them by value.
    • Return type. collect() needs to know what to build. The function's return type Vec<Vec<i32>> tells it.

    Complexity

    Time: O(n log k). Each of the n points is pushed once, and at most one entry is popped per point. The heap never holds more than k + 1 entries, so each push or pop costs O(log k). Sorting the final k entries costs O(k log k), which is no more than O(n log k) because k is at most n.

    Space: O(k). The heap holds at most k + 1 tuples, and the answer has k points.

    Edge cases and common mistakes

    • Ties in distance. In the third example, [0, 2], [2, 0] and [0, -2] are all at the same distance. Putting x and y in the tuple after dist orders them as the answer needs: [0, -2], [0, 2], [2, 0].
    • Duplicate points. [[1, 1], [1, 1], [5, 5]] with k = 2 gives [[1, 1], [1, 1]]. Each copy is its own entry in the heap.
    • Largest coordinates. [10000, 10000] has a squared distance of 200,000,000, which fits in an i32 with room to spare. No square root and no i64 are needed.
    • Using a square root. f64 distances are not needed, and floats cannot be the heap's key type directly because they have no total order in Rust.
    • Keeping the wrong end. The heap must drop its biggest entry. A min-heap would throw away the closest point instead.
    • Comparing heap.len() > k directly. len() is a usize and k is an i32, so Rust rejects the comparison until k is converted.
    • Forgetting to drop the distance. The answer lists only coordinates, so each tuple must become vec![x, y].

    Testing

    • points = [[1, 3], [-2, 2]], k = 1 gives [[-2, 2]].
    • points = [[3, 3], [5, -1], [-2, 4]], k = 2 gives [[3, 3], [-2, 4]].
    • points = [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]], k = 4 gives [[1, 1], [0, -2], [0, 2], [2, 0]].
    • points = [[1, 1], [1, 1], [5, 5]], k = 2 gives [[1, 1], [1, 1]].
    • points = [[10000, 10000], [-10000, -10000]], k = 2 gives [[-10000, -10000], [10000, 10000]].
    Keep going with Rust500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free