Start learning free

LESSON 390 OF 541 · HEAP

K Closest Points to Origin in Swift: O(n log k) with a Max-Heap

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 11
    Ctrl + Enter to check

    SWIFT · HEAP

    K Closest Points To Origin in Swift, explained

    Problem overview

    kClosest receives points, an array where each points[i] = [x, y] is a point on the plane, and an integer k. It returns the k points closest to the origin (0, 0), as an [[Int]]. The result is sorted by distance with the closest first. Points at the same distance come in order of x, and points with the same x in order of y. The answer is guaranteed to be unique: every point left out is strictly farther away than every point returned.

    For points = [[3, 3], [5, -1], [-2, 4]] and k = 2 the squared distances are 18, 26 and 20. The two closest points are [3, 3] and [-2, 4], in that order, so the answer is [[3, 3], [-2, 4]].

    For [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]] with k = 4, the point [1, 1] is closest. Three points are at distance 2 and come by x, then y: [0, -2], [0, 2], [2, 0]. The answer is [[1, 1], [0, -2], [0, 2], [2, 0]], and [3, 3] is left out.

    The task asks for O(n log k) time and O(k) extra space, where n is the number of points. Sorting all the points by distance would cost O(n log n).

    Approach: keep the k best in a max-heap

    No square root is needed. A point is closer than another exactly when its x * x + y * y is smaller, and that value is a whole number, so comparisons are exact.

    The idea is to keep the k closest points seen so far in a max-heap ordered by distance. A max-heap keeps its biggest item on top, so the top is always the farthest point among the ones kept: the one to drop when something better arrives. For each point:

    1. insert it into the heap;
    2. if the heap now holds more than k points, remove the top. Out of k + 1 points, the farthest cannot be among the k closest.

    After the last point, the heap holds the k closest. The heap order isn't sorted order, so the points are taken out and sorted by distance, then x, then y. The heap never holds more than k + 1 entries, so every insert and extract costs O(log k), which gives O(n log k) in total.

    The lesson provides a ready-to-use MaxHeap: MaxHeap(less) where less(a, b) is true when a is smaller, with insert(value), extract(), peek(), count and isEmpty. You don't write it yourself.

    Step by step

    1. Create the heap with var heap = MaxHeap<[Int]>({ $0[0] < $1[0] }). Each entry will be [dist, x, y], and the closure compares two entries by their first item, the distance. It is a var because inserting and extracting change it.
    2. Loop over the points with for point in points. Each point gets one chance to enter.
    3. Name the coordinates: let x = point[0], y = point[1].
    4. Compute the squared distance dist = x * x + y * y.
    5. Insert the entry [dist, x, y]. The entry carries the point along with the number the heap orders by.
    6. Trim: if heap.count > k, call heap.extract() to drop the farthest entry.
    7. After the loop, create closest as an empty [[Int]].
    8. Drain the heap: loop while !heap.isEmpty.
    9. Move entries: each turn, append heap.extract() to closest. That runs k times.
    10. Sort closest by the tuple of distance, x and y, taken from indexes 0, 1 and 2.
    11. Strip the distances with closest.map { [$0[1], $0[2]] } and return the result.

    The solution in Swift

    class Solution {
    func kClosest(_ points: [[Int]], _ k: Int) -> [[Int]] {
        var heap = MaxHeap<[Int]>({ $0[0] < $1[0] })
        for point in points {
            let x = point[0], y = point[1]
            let dist = x * x + y * y
            heap.insert([dist, x, y])
            if heap.count > k {
                heap.extract()
            }
        }
        var closest = [[Int]]()
        while !heap.isEmpty {
            closest.append(heap.extract())
        }
        closest.sort { ($0[0], $0[1], $0[2]) < ($1[0], $1[1], $1[2]) }
        return closest.map { [$0[1], $0[2]] }
    }
}

    Swift details worth knowing:

    • Closure shorthand. $0 and $1 are the closure's first and second arguments. { $0[0] < $1[0] } reads "entry one is smaller than entry two when its distance is smaller", which is the question a max-heap asks to decide what stays on top.
    • Tuples compare item by item. (a, b, c) < (d, e, f) looks at a versus d first, and only on a tie at b versus e, then c versus f. For example (1, -2, 2) < (1, 0, 2) is true. This gives the order the answer needs with no extra code.
    • Trailing closure. closest.sort { ... } passes the comparison closure after the parentheses. sort changes the array in place, which is why closest is a var; sorted would return a new array.
    • map builds a new array by applying the closure to each entry; here it turns [dist, x, y] into [x, y].
    • 64-bit Int. The largest squared distance with coordinates up to 10,000 is 200,000,000, far below the Int limit. In Swift, overflow stops the program instead of wrapping, but that can't happen here.

    Complexity

    Time: O(n log k). Each of the n points is inserted into a heap that holds at most k + 1 entries, and possibly followed by one extraction. Both cost O(log k). Draining the heap and sorting the k survivors cost O(k log k), which is no more than O(n log k) because k <= n.

    Space: O(k). The heap holds up to k + 1 entries of three numbers each, and closest holds k.

    Edge cases and common mistakes

    • A min-heap instead of a max-heap. The top must be the farthest kept point, since that is the one to discard. A min-heap would throw away the closest points.
    • Forgetting the final sort. A heap is only partly ordered, and extracting from it gives the farthest point first, so the entries still need sorting by distance, x and y.
    • Square roots. They aren't needed and would bring in floating-point rounding. Compare squared distances.
    • Duplicate points. The same point can appear twice. For [[1, 1], [1, 1], [2, 2]] with k = 2 the answer is [[1, 1], [1, 1]].
    • k equal to the number of points. Nothing is ever trimmed, and the function returns all the points sorted: [[3, 3], [5, -1], [-2, 4]] with k = 3 gives [[3, 3], [-2, 4], [5, -1]].
    • The origin itself. [[0, 0]] with k = 1 gives [[0, 0]]; its distance is 0.
    • Points on opposite sides. [10000, 10000] and [-10000, -10000] have the same squared distance, 200,000,000, so they are ordered by x: the answer for k = 2 is [[-10000, -10000], [10000, 10000]].

    Testing

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