Start learning free

LESSON 389 OF 540 · HEAP

K Closest Points to Origin in Go: 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 12
    Ctrl + Enter to check

    GO · HEAP

    K Closest Points To Origin in Go, explained

    Problem overview

    The function kClosest(points [][]int, k int) [][]int receives a slice of points, where points[i] = []int{x, y}, and a number k. It returns the k points closest to the origin (0, 0), 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 farther than every point kept.

    In the second example, points = [[3, 3], [5, -1], [-2, 4]] and k = 2. The squared distances are 18, 26 and 20, so the two closest are [3, 3] and [-2, 4], in that order. In the third example, [0, -2], [0, 2] and [2, 0] are all at the same distance, so they come ordered by x and then y: [1, 1], [0, -2], [0, 2], [2, 0].

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

    Approach: keep the k best in a max-heap

    Two ideas make this efficient.

    Compare squared distances. The distance of [x, y] is the square root of x*x + y*y, but the square root is increasing: one point is closer than another exactly when its squared distance is smaller. So the code never takes a square root, and the value stays a whole number. With coordinates up to 10,000 in absolute value, a squared distance is at most 200,000,000, far below what Go's 64-bit int holds.

    Keep only the k closest points seen so far. A max-heap keeps its biggest entry on top, so the top is always the farthest of the points kept. For each new point, push it; if the heap now holds k + 1 points, the farthest of them can't be among the k closest overall, so pop it. After the last point, the heap holds exactly the answer. Each push and pop costs O(log k) because the heap never holds more than k + 1 entries, which gives O(n log k).

    Each heap entry is a slice {dist, x, y}. The lesson provides the maxHeap type: its entries are []int, ordered by their first number, the biggest on top, so the distance decides.

    A heap is only partly ordered: knowing the top says little about the order of the rest. For the five numbers 5, 1, 4, 2, 3 pushed in that order, a max-heap of one-number entries holds them as [5 3 4 1 2] in memory. So at the end the k entries still need sorting, which costs only O(k log k).

    Here is the heap on the second example, printing its slice after each push:

    after push [[18 3 3]]
after push [[26 5 -1] [18 3 3]]
after push [[26 5 -1] [18 3 3] [20 -2 4]]
pop [26 5 -1]
[[20 -2 4] [18 3 3]]

    The third push makes three entries with k = 2, so the farthest, [26 5 -1], leaves.

    Step by step

    1. Create the heap. pointHeap := &maxHeap{} is a pointer to an empty maxHeap. The name heap is taken by the container/heap package, so the variable gets another name.
    2. Visit every point with for _, point := range points. Each point gets one chance to enter the heap.
    3. Name the coordinates: x, y := point[0], point[1].
    4. Compute the squared distance: dist := x*x + y*y.
    5. Push the entry []int{dist, x, y} with heap.Push(pointHeap, ...). The package function restores the heap order after the insert; calling the type's own Push method would only append to the slice.
    6. Trim to k. If pointHeap.Len() > k, call heap.Pop(pointHeap). The popped entry is the farthest one. After every point, the heap holds the k closest seen so far.
    7. Take the entries out of the heap: closest := *pointHeap is the slice the pointer points to.
    8. Sort them with slices.SortFunc(closest, slices.Compare). slices.Compare compares two slices element by element, and the first difference decides: distance first, then x, then y, which is the order the statement asks for.
    9. Allocate the result: result := make([][]int, k) has room for k points.
    10. Fill it with for i, entry := range closest.
    11. Drop the distance: result[i] = []int{entry[1], entry[2]} keeps only the coordinates.
    12. Return result.

    The solution in Go

    func kClosest(points [][]int, k int) [][]int {
    pointHeap := &maxHeap{}
    for _, point := range points {
        x, y := point[0], point[1]
        dist := x*x + y*y
        heap.Push(pointHeap, []int{dist, x, y})
        if pointHeap.Len() > k {
            heap.Pop(pointHeap)
        }
    }
    closest := *pointHeap
    slices.SortFunc(closest, slices.Compare)
    result := make([][]int, k)
    for i, entry := range closest {
        result[i] = []int{entry[1], entry[2]}
    }
    return result
}

    The Go details worth a closer look:

    • container/heap works on an interface. A type becomes a heap by providing Len, Less, Swap, Push and Pop. The lesson already provides that for maxHeap, so the solution only calls the package functions heap.Push and heap.Pop on a *maxHeap. They keep the heap order; calling the Push and Pop methods directly would not.
    • heap.Pop returns any. The solution ignores the popped value, so no type assertion is needed here.
    • *pointHeap. Dereferencing the pointer gives the slice of entries itself. slices.SortFunc then sorts that slice in place, and the heap order no longer matters because nothing pushes to the heap again.
    • slices.SortFunc(closest, slices.Compare). The comparison function must return a negative number, zero or a positive number. slices.Compare does exactly that for slices of ordered elements, so it can be passed as it is. For example, comparing [8 3 3] with [8 -2 4] gives 1, because 3 is greater than -2 at the first difference.
    • Sorting needs no tie rule of its own. Entries with the same distance are ordered by x then y because those numbers follow the distance in each entry.
    • make([][]int, k). It creates k nil slices, each replaced in the loop. The loop assumes the heap holds exactly k entries, which is true because k <= len(points).

    Complexity

    Time: O(n log k). Each of the n points costs one push and at most one pop on a heap of at most k + 1 entries, which is O(log k) each. Sorting the final k entries adds O(k log k), which is no bigger. Sorting all n points would cost O(n log n), which is worse when k is small.

    Space: O(k). The heap holds at most k + 1 entries of three integers, and the result holds k points.

    Edge cases and common mistakes

    • The origin itself. [[0, 0]] with k = 1 gives [[0, 0]], with distance 0.
    • Duplicate points. [[1, 1], [1, 1], [2, 2]] with k = 2 gives [[1, 1], [1, 1]]. Equal entries are both kept.
    • k equal to the number of points. Nothing is ever popped, and the function returns every point sorted. For the second example with k = 3 the result is [[3, 3], [-2, 4], [5, -1]].
    • A min-heap instead of a max-heap. A min-heap keeps the closest point on top, which is the opposite of what is needed. The top has to be the farthest kept point, the one to drop.
    • Returning in heap order. Without the final sort, the points come out in the heap's internal order, not by distance. The extra O(k log k) sort is needed.
    • Taking square roots. They are unnecessary, and floating-point values add rounding questions. The integer x*x + y*y is exact.
    • Negative coordinates. Squaring removes the sign, so [-2, 2] and [2, 2] are equally far. They differ only in the tie-break by x.
    • Largest values. [10000, 10000] has a squared distance of 200,000,000, which fits in an int without overflow.
    • Calling the methods instead of the package functions. pointHeap.Push(...) appends without restoring the heap order, and the top is no longer the farthest point.

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