Start learning free

LESSON 389 OF 540 · HEAP

K Closest Points to Origin in Kotlin: 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

    KOTLIN · HEAP

    K Closest Points To Origin in Kotlin, explained

    Problem overview

    The function kClosest(points: Array<IntArray>, k: Int): Array<IntArray> receives points [x, y] on a plane and returns the k points closest to the origin (0, 0). The answer is sorted: closest first, points at the same distance ordered by x, and points with the same x ordered by y. The answer is guaranteed to be unique, so every point left out is strictly farther away than every point in the answer.

    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 result is [[3, 3], [-2, 4]].

    For [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]] and k = 4 the result is [[1, 1], [0, -2], [0, 2], [2, 0]]. Three points are at distance 2, so they are ordered by x, and the two with x = 0 by y.

    The lesson asks for O(n log k) time and O(k) extra space. Sorting all n points would cost O(n log n).

    Approach: a heap that keeps only the k closest

    Two facts shape the solution.

    Squared distances are enough. 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 Int.

    Only the k closest need to be remembered. Scan the points once and keep a collection of the best k seen so far. When a new point arrives, it joins the collection, and if the collection now holds k + 1 points, the farthest of them is certainly not among the final k, so it is dropped. The collection must therefore answer "which point is the farthest?" quickly, and that is the job of a max-heap: the biggest distance sits on top, and adding or removing costs O(log k) because the heap never holds more than k + 1 entries.

    For example, with k = 2 on the points [3, 3], [5, -1], [-2, 4], each entry written as [distance, x, y]:

    after (3, 3): kept=[[18, 3, 3]]
after (5, -1): kept=[[18, 3, 3], [26, 5, -1]]
after (-2, 4): dropped=[26, 5, -1] kept=[[18, 3, 3], [20, -2, 4]]

    When the third point arrives, the heap has three entries, and the farthest, 26, is removed. The scan costs n heap operations of O(log k) each, which gives O(n log k).

    A heap is only partly ordered, so the k survivors still have to be sorted at the end. That costs O(k log k), which is no more than the rest of the work since k <= n.

    Step by step

    1. Create the heap. It is a PriorityQueue<IntArray> built with compareByDescending { it[0] }. Each entry is intArrayOf(dist, x, y), and the comparator looks at the distance, so the biggest one is on top.
    2. Loop over the points, taking each one apart as (x, y).
    3. Compute the squared distance dist as x * x + y * y.
    4. Add the entry intArrayOf(dist, x, y). The distance is stored with the point because it is what the heap orders by.
    5. Trim the heap. If heap.size > k, call heap.poll(). It removes the top, the farthest point. After each point, the heap holds the k closest seen so far.
    6. Copy the survivors into a list with heap.toMutableList(). The list is in heap order, not sorted.
    7. Sort the list with three keys in order: distance, then x, then y. Each later key matters only when the earlier ones are equal.
    8. Drop the distances by mapping each entry to intArrayOf(it[1], it[2]), then convert the list to an array with toTypedArray().

    The solution in Kotlin

    class Solution {
    fun kClosest(points: Array<IntArray>, k: Int): Array<IntArray> {
        val heap = PriorityQueue<IntArray>(compareByDescending { it[0] })
        for ((x, y) in points) {
            val dist = x * x + y * y
            heap.add(intArrayOf(dist, x, y))
            if (heap.size > k) {
                heap.poll()
            }
        }
        val closest = heap.toMutableList()
        closest.sortWith(compareBy<IntArray>({ it[0] }, { it[1] }, { it[2] }))
        return closest.map { intArrayOf(it[1], it[2]) }.toTypedArray()
    }
}

    Kotlin details worth a look:

    • PriorityQueue is Java's. Kotlin has no heap of its own, so it uses java.util.PriorityQueue, which is a min-heap: the smallest element by the comparator is on top. The comparator compareByDescending { it[0] } reverses the order on the distance, which turns it into a max-heap. In a file of your own, add import java.util.PriorityQueue.
    • Destructuring an IntArray. for ((x, y) in points) uses the array's component1() and component2() functions, which return its first two elements. A point with more than two numbers would simply have the rest ignored.
    • add and poll. add inserts an element. poll removes and returns the top, or null when the queue is empty. Here the heap is never empty when poll runs, and the result isn't needed.
    • The heap is not sorted. Iterating a PriorityQueue, which is what toMutableList() does, doesn't give sorted order: for entries with distances 1 to 6 added in order, the list comes out as [6, 4, 5, 1, 3, 2]. Only the top is guaranteed. This is why the final sort is needed.
    • compareBy with several selectors. compareBy<IntArray>({ it[0] }, { it[1] }, { it[2] }) compares by the first selector and goes to the next only on a tie. The explicit type argument <IntArray> is needed because the lambdas alone don't tell the compiler what they receive.
    • Converting back. map on a list returns a List, and the function must return an Array<IntArray>, so toTypedArray() finishes the job.
    • Why the answer's uniqueness matters. The heap decides which point to drop by distance alone. If two points tied at the cut-off, it could drop either one. The lesson guarantees there is no such tie.

    Complexity

    Time: O(n log k). Each of the n points is added once and the heap is trimmed when it passes k entries, so it never holds more than k + 1 entries and each operation costs O(log k). The final sort of k entries costs O(k log k).

    Space: O(k) extra. The heap, the list copy and the returned array all hold at most k entries.

    Edge cases and common mistakes

    • The origin itself. [[0, 0]] with k = 1 returns [[0, 0]]. Its distance is 0, which is a valid value.
    • Duplicate points. [[1, 1], [1, 1], [5, 5]] with k = 2 returns [[1, 1], [1, 1]]. Both copies are kept, since each is its own entry in the heap.
    • Largest coordinates. [[10000, 10000], [-10000, -10000], [0, 1]] with k = 3 returns [[0, 1], [-10000, -10000], [10000, 10000]]. The two far points are at the same distance, so the one with the smaller x comes first, and x * x + y * y reaches 200,000,000, which still fits in an Int.
    • Using a min-heap by mistake. With the farthest point not on top, poll would drop the closest point instead.
    • Skipping the final sort. The heap's iteration order isn't distance order, so the tests would fail.
    • Forgetting the tie-breakers. Sorting only by distance leaves points at equal distance in an arbitrary order.
    • Using square roots. They add nothing here. Comparing squared distances gives the same order without floating-point values.

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