Start learning free

LESSON 287 OF 414 · DIVIDE AND CONQUER

Smallest K in Scala: Quickselect in O(n) Expected Time

Practice O(n) 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:

Your algorithm should run in O(n) time on average and use O(log n) extra space on average (for the recursion), when partitioning in place.


STEP 1 OF 18
    Ctrl + Enter to check

    SCALA · DIVIDE AND CONQUER

    Smallest K in Scala, explained

    Problem overview

    The method kSmallest receives an Array[Int] called nums and an integer k. It returns the k smallest values, in any order. For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4 the answer holds 1, 2, 2 and 3; [2, 1, 3, 2] is just as correct as [1, 2, 2, 3].

    The target is O(n) time and O(log n) extra space.

    Approach: partition, then follow only the side that matters

    Sorting the whole array and taking the first k values costs O(n log n), and keeping a heap of size k costs O(n log k). This lesson goes further, using the idea behind Quick Sort. It is known as quickselect.

    A partition picks a pivot and rearranges a range so that the values less than or equal to the pivot come first, then the pivot, then the larger values. The pivot lands at its final sorted index pivot. That means the range [left, pivot] holds the pivot - left + 1 smallest values of the range. Compare that count with k:

    • If the count is at least k, all k wanted values are inside [left, pivot]. The pivot is the largest value there, so keep searching in [left, pivot - 1] with the same k.
    • If the count is smaller, every value in [left, pivot] belongs to the answer, and only k - count more are needed from [pivot + 1, right].

    Unlike Quick Sort, it recurses into one side only. With random pivots the range shrinks by a constant factor on average, so the total work is about n + n/2 + n/4 + ..., which is O(n).

    Step by step

    1. Declare smallestKNums(nums, left, right, k) returning Unit. It rearranges nums in place so that the k smallest values of the range come first.
    2. Base case: if left >= right, there is at most one value and nothing to rearrange, so return.
    3. Declare partition(nums, left, right) returning the pivot's final index.
    4. Pick a random pivot index: left plus (scala.util.Random.nextDouble() * (right - left + 1)).toInt. nextDouble() returns a number from 0 up to but not including 1, so the product, truncated, is a whole number from 0 to right - left.
    5. Write swap (above partition) to exchange two array elements through a temporary val.
    6. Move the pivot to the end by swapping randomIndex with right.
    7. Set up the scan: val value = nums(right) is the pivot, and var correctIndex = left is where the next value that is at most value goes.
    8. Scan with a while loop over i from left to right inclusive. The pivot, at right, is visited last.
    9. Swap qualifying values forward. If nums(i) <= value, swap nums(i) with nums(correctIndex) and add one to correctIndex.
    10. Return correctIndex - 1. The pivot was the last value moved, so it sits just before correctIndex.
    11. Back in smallestKNums, call partition and store the result as pivot.
    12. Branch on the count pivot - left + 1 >= k: recurse on [left, pivot - 1] with the same k.
    13. Otherwise recurse on [pivot + 1, right] with k - (pivot - left + 1).
    14. In kSmallest, call smallestKNums(nums, 0, nums.length - 1, k).
    15. Return nums.slice(0, k), the first k positions, as a new array.

    The solution in Scala

    class Solution {
    def swap(nums: Array[Int], i: Int, j: Int): Unit = {
        val temp = nums(i)
        nums(i) = nums(j)
        nums(j) = temp
    }

    def partition(nums: Array[Int], left: Int, right: Int): Int = {
        val randomIndex = left + (scala.util.Random.nextDouble() * (right - left + 1)).toInt
        swap(nums, randomIndex, right)

        val value = nums(right)
        var correctIndex = left
        var i = left
        while (i <= right) {
            if (nums(i) <= value) {
                swap(nums, i, correctIndex)
                correctIndex += 1
            }
            i += 1
        }
        correctIndex - 1
    }

    def smallestKNums(nums: Array[Int], left: Int, right: Int, k: Int): Unit = {
        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))
        }
    }

    def kSmallest(nums: Array[Int], k: Int): Array[Int] = {
        smallestKNums(nums, 0, nums.length - 1, k)
        nums.slice(0, k)
    }
}

    Scala details:

    • while instead of for. The loop uses var i and i += 1 because Scala has no ++. A for (i <- left to right) would also work.
    • Random numbers. scala.util.Random.nextInt(n) gives a number from 0 to n - 1 directly. The lesson multiplies nextDouble() instead, and .toInt truncates the fraction.
    • slice(from, until) excludes the until index and returns a new array, so the result is a copy.
    • The input is rearranged. After the call, the caller's nums is in a different order. For the example it becomes 2,1,2,3,3,5,5,6,7,9 (the exact order depends on the random pivots).
    • Result of a method. partition and kSmallest end with an expression instead of return. The early return in smallestKNums works because its result type is declared.

    Complexity

    Time: O(n) expected. The first partition scans n values. A random pivot lands in the middle part of the range often enough that the range we continue with is a constant fraction smaller on average, so the scans add up to a constant multiple of n. The worst case, with consistently bad pivots, is O(n²).

    Space: O(log n) expected for the recursion, not counting the k-element result array.

    Edge cases and common mistakes

    • k equal to the length. For [4, 4, 1] and k = 3 the answer is all three values, 1, 4, 4.
    • k larger than the length. The search runs out of values, and slice stops at the end of the array, so you get every element.
    • k = 0. The result has length 0.
    • Single element or empty array. [7] with k = 1 returns [7], and an empty array with k = 0 returns an empty array.
    • Duplicates. Equal values are fine, as in the example with two 2s. But if every value is the same, the test <= sends them all to the left, the pivot ends at right, and each call shrinks the range by one element. With 1,000 equal values and k = 1 the code calls partition 999 times, which is quadratic work.
    • Already sorted input. Because the pivot is random, it stays fast. For 1,000 sorted values and k = 1 it needs only a handful of calls to partition, and the exact count varies from run to run.
    • Using the wrong k in the second branch. After keeping count values, only k - count remain to be found. Passing the original k asks for too many.
    • Expecting sorted output. Only the set of k values is guaranteed, not their order.

    Testing

    • [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] with k = 4 gives the values 1, 2, 2, 3 in some order.
    • [4, 4, 1] with k = 3 gives 1, 4, 4 in some order.
    • [7] with k = 1 gives [7].
    • [4, 4, 1] with k = 0 gives an empty array.
    Keep going with ScalaAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free