Start learning free

LESSON 292 OF 348 · HEAP

Smallest K Integers in Kotlin: O(n log k) with a Max Heap

Practice O(n log k) Read the full explanation

Given an array of positive integers nums, return the smallest k values, in any order you want.

Example:

Input: nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4
Output: [1, 2, 2, 3]
Explanation: Smallest number is 1, 2nd smallest is 2, 
            3rd smallest is 2, 4th smallest is 3

The result can be in any order, [2, 1, 3, 2] is also a correct answer.

Note:

For this lesson, your algorithm should run in O(n log k) time and use O(k) extra space.
(There are faster solutions which we will discuss in future lessons)


STEP 1 OF 6
    Ctrl + Enter to check

    KOTLIN · HEAP

    Smallest K in Kotlin, explained

    Problem overview

    Given an IntArray of positive integers nums and a number k, return the k smallest values, in any order. Duplicates count separately: with nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4 the smallest values are 1, 2, 2 and 3, because the array holds two 2s and two 3s and only one of the 3s makes the cut.

    The lesson asks for O(n log k) time and O(k) extra space. It says that faster solutions exist and are covered in later lessons.

    Approach

    The simplest solution sorts the whole array and takes the first k values. That costs O(n log n) time. It works, but it does more than the problem needs: it puts all n values in order when only the k smallest are wanted.

    A heap does better by remembering only k values at a time. The idea:

    • Keep the k smallest values seen so far in a max heap, a heap whose top is the greatest value it holds.
    • That top is the weakest member of the group, the first one that should be thrown out.
    • When a new value arrives and the heap is full, compare it with the top. If the new value is smaller, it belongs in the group and the top does not, so swap them. Otherwise ignore the new value.

    The heap never holds more than k values, so each add or remove costs O(log k). With n values to process, the total is O(n log k), and the heap itself takes O(k) space.

    Step by step

    1. Build a max heap. Kotlin on the JVM uses Java's java.util.PriorityQueue, which is a min heap by default: the smallest value comes out first. Passing the comparator compareByDescending { it } reverses the order, so the greatest value is on top. The lambda picks the thing to compare, here the number itself.
    2. Visit every value. A for loop with the variable value goes over nums.
    3. Fill the heap first. While the heap has fewer than k values (size < k), every value is accepted with add.
    4. Compare with the top. Once the heap is full, peek() reads the greatest kept value without removing it. Only a value strictly smaller than that top deserves a place.
    5. Swap. In that case, poll() removes the current top and add inserts the new value. The heap keeps exactly k values.
    6. Empty the heap into an array. After the loop the heap holds the k smallest values. Create result as an IntArray(k) and fill positions 0 until k with poll() calls. Because the top is always the greatest, the values come out from the largest to the smallest. The problem accepts any order, so that's fine.

    The solution in Kotlin

    class Solution {
    fun kSmallest(nums: IntArray, k: Int): IntArray {
        val smallestKNums = java.util.PriorityQueue<Int>(compareByDescending { it })

        for (value in nums) {
            if (smallestKNums.size < k) {
                smallestKNums.add(value)
            } else if (value < smallestKNums.peek()) {
                smallestKNums.poll()
                smallestKNums.add(value)
            }
        }

        val result = IntArray(k)
        for (i in 0 until k) {
            result[i] = smallestKNums.poll()
        }
        return result
    }
}

    Kotlin details worth knowing:

    • The class is written with its full name, java.util.PriorityQueue, so no import is needed. You could write import java.util.PriorityQueue at the top of the file and use the short name instead.
    • compareByDescending { it } returns a Comparator. Kotlin works out the type of it from PriorityQueue<Int>, so the lambda needs no annotations.
    • Methods that come from Java, such as peek() and poll(), return a "platform type" (Int!). Kotlin lets you use the result as an Int without a null check. In real Java code peek() and poll() return null on an empty queue, which would crash on the next use as an Int. This solution never calls them on an empty heap, because peek() runs only after the heap has reached k values. That is safe as long as k is at least 1.
    • 0 until k excludes k, so it visits exactly the indices 0 to k - 1 of the result.
    • IntArray(k) creates an array of k zeros, which the loop then overwrites.

    Complexity

    Time: O(n log k). Each of the n values costs at most one peek (constant time) plus one poll and one add, each O(log k) because the heap never exceeds k values. Emptying the heap at the end is k polls, which is O(k log k) and never more than the loop's cost.

    Space: O(k). The heap holds at most k values. The result array also has size k.

    When k is much smaller than n, this beats the O(n log n) sort. When k is close to n, the two are about equal.

    Edge cases and common mistakes

    • Using a min heap. A min heap keeps the smallest value on top, so you would throw away the wrong value. The comparator is what turns the queue into a max heap.
    • Duplicates. The comparison is value < peek(), strictly smaller. A value equal to the top changes nothing in the result, so skipping it is correct.
    • Calling peek() too early. On an empty queue it returns null, which Kotlin would turn into a crash when treated as an Int. The size < k branch comes first so that the heap is full before peek() runs.
    • k equal to the array size. Every value is added and none is ever replaced. The result is the whole array, in descending order.
    • Expecting sorted output. The result comes out in descending order here, so for the example it is [3, 2, 2, 1]. Any order is accepted.

    Testing

    • [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] with k = 4 returns [3, 2, 2, 1].
    • [4, 1] with k = 2 returns [4, 1], the whole array in descending order.
    Keep going with KotlinAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free