LESSON 263 OF 348 · DIVIDE AND CONQUER

Smallest K Integers in Kotlin: O(n) Quickselect

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:

Your algorithm should run in O(n) time and use O(log n) extra space.

STEP 1 OF 15

Guided

Ctrl + Enter to check

KOTLIN · DIVIDE AND CONQUER

Smallest K in Kotlin, explained

Problem overview

Given an IntArray of positive integers nums and a number k, kSmallest returns the k smallest values, in any order. For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4 the smallest values are 1, 2, 2 and 3. Duplicates count separately, so the answer is [1, 2, 2, 3] and [2, 1, 3, 2] is just as good.

The lesson asks for O(n) time and O(log n) extra space.

Approach

The obvious solution sorts the array and takes the first k values. That costs O(n log n), and it does more than the task needs: we don't care about the order of the answer, and we don't care about the order of the values that are left out.

Quickselect does less work by reusing the partition step of Quick Sort. After partitioning around a pivot, the pivot sits at its final sorted index, with smaller-or-equal values before it and larger ones after it. That tells us how many values are at most the pivot: pivot - left + 1, counting the pivot itself. Compare that count with k:

If the count is at least k, all k smallest values lie in [left, pivot]. Only that part needs more work, and the pivot, being the largest value there, can stay put. Recurse on [left, pivot - 1] with the same k.
If the count is smaller, everything in [left, pivot] belongs to the answer already. Recurse on the right part only, looking for the remaining k - (pivot - left + 1) values.

Unlike Quick Sort, which recurses on both sides, Quickselect follows one side. With a random pivot the range shrinks by a constant factor on average, so the work is n + n/2 + n/4 + ... = O(n) in expectation. At the end, the first k positions of the array hold the answer.

Step by step

Declare smallestKNums with the array, the bounds left and right and k. It returns nothing; it rearranges nums so that the k smallest values of the range come first.
Base case. If left >= right, the range holds at most one value. Return.
Declare partition. It returns the pivot's final index, after moving the values that are at most the pivot in front of it and the larger ones behind it.
Pick the pivot randomly. left + (Math.random() * (right - left + 1)).toInt() is a random index in left..right. Math.random() gives a double from 0 up to but not including 1, and toInt() drops the fraction.
Write swap with a temporary variable.
Move the pivot to the end of the range with swap.
Keep two variables. val value = nums[right] is the pivot, and var correctIndex = left is where the next value that is no larger than the pivot goes.
Scan i over left..right, both ends included, so the pivot is visited last.
Move small values forward. If nums[i] <= value, swap it with nums[correctIndex] and increment correctIndex.
Return correctIndex - 1. The pivot was the last value moved, so it ends just before correctIndex.
Use the pivot in smallestKNums. After the base case, store the result of partition in pivot.
Left branch. If pivot - left + 1 >= k, recurse on [left, pivot - 1] with the same k. When the count equals k exactly the call has nothing left to find, but it only rearranges values inside [left, pivot - 1], so it is harmless.
Right branch. Otherwise, recurse on [pivot + 1, right] with k - (pivot - left + 1).
Start the search in kSmallest on 0 to nums.size - 1.
Return a copy of the first k values with nums.copyOfRange(0, k). The end index is exclusive.

The solution in Kotlin

class Solution {
    fun swap(nums: IntArray, i: Int, j: Int) {
        val temp = nums[i]
        nums[i] = nums[j]
        nums[j] = temp
    }

    fun partition(nums: IntArray, left: Int, right: Int): Int {
        val randomIndex = left + (Math.random() * (right - left + 1)).toInt()
        swap(nums, randomIndex, right)
        val value = nums[right]
        var correctIndex = left
        for (i in left..right) {
            if (nums[i] <= value) {
                swap(nums, i, correctIndex)
                correctIndex++
            }
        }
        return correctIndex - 1
    }

    fun smallestKNums(nums: IntArray, left: Int, right: Int, k: Int) {
        if (left >= right) {
            return
        }
        val pivot = partition(nums, left, right)
        if (pivot - left + 1 >= k) {
            smallestKNums(nums, left, pivot - 1, k)
        } else {
            smallestKNums(nums, pivot + 1, right, k - (pivot - left + 1))
        }
    }

    fun kSmallest(nums: IntArray, k: Int): IntArray {
        smallestKNums(nums, 0, nums.size - 1, k)
        return nums.copyOfRange(0, k)
    }
}

Kotlin details:

The input is rearranged. nums is an array, passed by reference, so after the call it is no longer in its original order. For the example, a run can leave it as [1, 2, 2, 3, 3, 9, 5, 6, 7, 5], though the order of the tail varies from run to run because the pivot is random. The returned array is a separate copy, so it isn't affected by later changes to nums.
copyOfRange(from, to) includes from and excludes to, and it throws IndexOutOfBoundsException if to is larger than the array size.
Math.random() is Java's static method, available on the JVM. Multiplying by the range size and calling toInt() truncates, giving an index from left to right. The Kotlin form (left..right).random() does the same.
val versus var. value never changes, so it's a val. correctIndex is updated and must be a var.

Complexity

Time. Each call partitions its range in O(range). With a random pivot, the expected range after one step is a constant fraction of the previous one, so the sum is O(n) in expectation. The worst case, with consistently bad pivots, is O(n²), which random choice makes very unlikely.

Space. The partition works in place and there is only one recursive call per level. The stack is O(log n) deep in expectation. The returned array adds O(k), and it is the answer itself, not working space.

Edge cases and common mistakes

k equal to the array size. The result holds every value, in some order.
k = 1. The search closes in on the minimum.
k = 0. The left branch is always taken and the range shrinks until the base case. The result is an empty array.
Duplicates. Values equal to the pivot count toward the left side, which is what makes [4, 4, 4] with k = 2 give [4, 4].
Using pivot in the recursive call on the left. Going to pivot - 1 avoids reconsidering a value that is already in place.
Forgetting to subtract in the right branch. If you keep the old k, the search looks for too many values.
Expecting a sorted result. The output order is not guaranteed. Sort the returned array yourself if a test needs it.

Testing

kSmallest([5, 9, 3, 6, 2, 1, 3, 2, 7, 5], 4) returns 1, 2, 2 and 3 in some order.
kSmallest([4, 4, 4], 2) returns [4, 4].
kSmallest([8, 2, 5], 1) returns [2].
kSmallest([8, 2, 5], 3) returns 2, 5 and 8 in some 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