Start learning free

LESSON 223 OF 414 · SORTING

Smallest K in Scala: O(n log n) with Sorting

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

For this lesson, your algorithm should run in O(n log n) time and use O(1) extra space besides the returned values, when sorting in place (sorting a copy of nums takes O(n) extra space).
(There are faster solutions which we will discuss in future lessons)


STEP 1 OF 2
    Ctrl + Enter to check

    SCALA · SORTING

    Smallest K in Scala, explained

    Problem overview

    The method kSmallest receives an array nums of positive integers and a number 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 sorted array is [1, 2, 2, 3, 3, 5, 5, 6, 7, 9], so the four smallest values are 1, 2, 2, 3. Duplicates count separately: the second 2 is the third smallest value. The order of the answer doesn't matter, so [2, 1, 3, 2] would be accepted too.

    Approach: sort, then take a prefix

    If the array is in ascending order, the k smallest values are simply the first k elements. Sorting turns "find the smallest k" into "cut off the front of an array". Sorting costs O(n log n) and slicing costs O(k), so the total is O(n log n), which is the budget the statement sets.

    The statement also mentions O(1) extra space. Sorting a copy of the array, as the reference solution does, uses O(n) extra memory, so it only fits that budget loosely. Sorting nums in place with scala.util.Sorting.quickSort(nums) and slicing afterwards would avoid the copy. Faster methods that avoid sorting everything exist, and the lesson leaves them for later.

    Step by step

    1. Sort the numbers. Calling sorted on an array returns a new array in ascending order and leaves nums as it was. Store it in sortedNums.
    2. Return the first k elements. slice(0, k) takes the elements from index 0 up to, but not including, index k. That is exactly k elements.

    The solution in Scala

    class Solution {
    def kSmallest(nums: Array[Int], k: Int): Array[Int] = {
        val sortedNums = nums.sorted
        return sortedNums.slice(0, k)
    }
}

    The Scala details worth knowing:

    • sorted is non-destructive. After the call, nums still holds its original order, and sortedNums is a separate array. Use val because the name is never reassigned.
    • slice(from, until) excludes the until index. sortedNums.take(k) gives the same result and reads a little more directly.
    • slice is forgiving about bounds. If k is larger than the array, you get the whole array back instead of an exception. With k = 0 you get an empty array.
    • return is optional. The last expression of a method is its value, so sortedNums.slice(0, k) alone would work. The lesson's code uses return, which is also valid here, because it is directly in the method body.
    • Arrays print as references. To see the result, use result.mkString(", "). Calling println on the array itself shows an object reference.

    Complexity

    Time: O(n log n). The sort dominates. The slice copies k elements, and k is at most n.

    Space: O(n) with this solution, because sorted allocates a second array, plus O(k) for the result. Sorting in place would bring the extra space down to the result alone.

    Edge cases and common mistakes

    • Duplicates. They are not merged. For [2, 2, 5] and k = 2 the answer is [2, 2].
    • k larger than the array. [3, 1] with k = 5 returns [1, 3], both elements.
    • k = 0. The result is an empty array.
    • Taking from the wrong end. The first k elements of the ascending array are the smallest. The last k are the largest.
    • Expecting nums to be sorted afterwards. It isn't. Only sortedNums is.
    • slice(0, k - 1). This is off by one and returns only k - 1 elements.

    Testing

    • nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives [1, 2, 2, 3].
    • nums = [3, 1], k = 5 gives [1, 3].
    • nums = [3, 1], 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