Start learning free

LESSON 389 OF 540 · HEAP

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

    SCALA · HEAP

    K Closest Points To Origin in Scala, explained

    Problem overview

    The method kClosest receives points, an Array[Array[Int]] where each point is [x, y], and an integer 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 points with the same x by y. The answer is unique: every point left out is farther away than every point kept.

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

    Approach: keep the k best in a max-heap

    Sorting all the points by distance and taking the first k works, but costs O(n log n). The statement asks for O(n log k) time and O(k) extra space, so we keep only the k closest points seen so far.

    The right tool is a max-heap: a structure where the biggest element is always on top and adding or removing it costs O(log size). Scala's PriorityQueue is a max-heap. If the entries are ordered by distance, the top is the farthest point we are keeping, which is exactly the one to drop when a closer point arrives. Add every point, and whenever the heap holds more than k entries, remove the top. After the last point, the heap holds the k closest.

    No square roots are needed. One point is closer than another exactly when its x * x + y * y is smaller, and this squared distance is a whole number. With coordinates up to 10,000 in absolute value it is at most 200,000,000, which fits in an Int.

    The heap never holds more than k + 1 entries, so each operation costs O(log k).

    Step by step

    1. Create the heap: an empty PriorityQueue[(Int, Int, Int)](). Each entry is the tuple (dist, x, y).
    2. Visit every point with for (point <- points). Each point gets one chance to enter the heap.
    3. Name the coordinates: x is point(0) and y is point(1).
    4. Compute the squared distance dist = x * x + y * y.
    5. Enqueue the tuple (dist, x, y). The outer parentheses belong to the call, the inner ones build the tuple.
    6. Trim the heap: if heap.size > k, call heap.dequeue(). Out of k + 1 points, the farthest can't be among the k closest.
    7. Sort what is left: heap.toArray.sorted copies the entries into an array and sorts it. Tuples sort item by item, so the order is distance, then x, then y.
    8. Drop the distances: closest.map { case (dist, x, y) => Array(x, y) } turns each tuple back into a point, and is the method's last expression.

    The solution in Scala

    In a file of your own, the code needs import scala.collection.mutable.PriorityQueue at the top.

    class Solution {
    def kClosest(points: Array[Array[Int]], k: Int): Array[Array[Int]] = {
        val heap = PriorityQueue[(Int, Int, Int)]()
        for (point <- points) {
            val x = point(0)
            val y = point(1)
            val dist = x * x + y * y
            heap.enqueue((dist, x, y))
            if (heap.size > k) {
                heap.dequeue()
            }
        }
        val closest = heap.toArray.sorted
        closest.map { case (dist, x, y) => Array(x, y) }
    }
}

    The Scala details worth a closer look:

    • Tuples compare item by item. (18, 3, 3) is smaller than (20, -2, 4) because 18 < 20; with equal distances the x decides, then the y. Scala provides this ordering for tuples of numbers, so PriorityQueue[(Int, Int, Int)]() needs no comparator.
    • Max-heap by default. dequeue() removes the biggest entry, which is what we want here. A min-heap would need a reversed ordering.
    • The heap's internal order isn't sorted. Listing the heap, for example with toList, gives the entries in its internal order. In the example, the heap holds (18, 3, 3) and (20, -2, 4), but toList gives (20, -2, 4) first. That is why the method sorts before returning.
    • Two sets of parentheses in enqueue((dist, x, y)). Without the inner pair, the three values would look like three arguments.
    • case (dist, x, y) => … takes each tuple apart into its three items. dist is not used on the right-hand side, which is fine.
    • The result type. map on an Array[(Int, Int, Int)] returns an Array[Array[Int]], which matches the method's return type.

    Complexity

    Time: O(n log k). Each of the n points is added once and the heap is trimmed at most once per point; the heap never holds more than k + 1 entries, so each operation costs O(log k). The final sort of k entries costs O(k log k), which is no more than O(n log k) because k <= n.

    Space: O(k). The heap, and then the sorted array, hold k entries.

    Edge cases and common mistakes

    • Ties in distance. For the points [0, 2], [2, 0], [1, 1], [0, -2], [3, 3] with k = 4, three points are at squared distance 4. They come out as [0, -2], [0, 2], [2, 0], by x and then y, after [1, 1]. The answer is [[1, 1], [0, -2], [0, 2], [2, 0]].
    • Repeated points. The same point may appear more than once. [[1, 1], [1, 1], [2, 2]] with k = 2 gives [[1, 1], [1, 1]].
    • k equals the number of points. Nothing is ever dequeued, and the method just sorts everything.
    • Storing only the distance. A heap of plain Int distances can't give the points back. The tuple carries x and y along, and also provides the tie-break.
    • Square roots. They aren't needed, and they would bring Double rounding into comparisons that are exact with integers.
    • Overflow. For the corner [10000, 10000] the squared distance is 200,000,000, comfortably below the Int limit of about 2.1 billion.

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