LESSON 255 OF 348 · BINARY SEARCH

Smallest K in Kotlin: O(n log n) with Binary Search on Values

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.

STEP 1 OF 3

Guided

Ctrl + Enter to check

KOTLIN · BINARY SEARCH

Smallest K in Kotlin, explained

Problem overview

The function kSmallest receives an IntArray of positive integers called nums and an Int called k. It returns an IntArray with 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, so [1, 2, 2, 3] is correct, and so is [2, 1, 3, 2].

The statement allows O(n log n) time. Sorting the array and copying the first k elements would also fit, but this lesson practises a different idea: binary search over the values themselves.

Approach: find the k-th smallest number, then collect

The plan has two parts.

Find the value of the k-th smallest number, call it kthNum, with a binary search. The search runs over possible values, not over indices: the array is not sorted, so there are no indices to cut in half. The candidates are every integer between the minimum and the maximum of nums.
Collect the answer. Everything strictly smaller than kthNum belongs to the result. If fewer than k numbers are strictly smaller, the rest of the result is filled with copies of kthNum.

For the first part, the test for a candidate value is "how many numbers in nums are less than value?". Call that count smaller. If smaller < k, then fewer than k numbers sit below value, so the k-th smallest is value or something larger. That makes value a possible answer: remember it and try larger values. If smaller >= k, at least k numbers are below value, so the k-th smallest is strictly less than value and the search goes lower.

The largest value that passes the test is the k-th smallest number. The count never decreases as value grows, which is what makes binary search legal.

On the example, the minimum is 1 and the maximum is 9. The middle value 5 has five numbers below it (3, 2, 1, 3, 2), which is at least 4, so the search goes lower. The middle of 1 to 4 is 2, with one number below it, so 2 is remembered and the search goes up. Then 3 has three numbers below it, still fewer than 4, so 3 is remembered. Next, 4 has five below it, which is too many. The remembered answer is 3. The numbers below 3 are 2, 1 and 2, and one copy of 3 completes the result: [2, 1, 2, 3].

Step by step

Write howManyLess(nums, value). Start a var counter at 0, loop over nums with num, add one each time value > num, and return the counter. The parameter is value because val is a Kotlin keyword.
Write kSmallestNum(nums, k). Start left and right at nums[0], then loop with minOf and maxOf so they end as the minimum and the maximum.
Remember the best candidate in var result, starting at left.
Run the search while left <= right. Take value = left + (right - left) / 2 and smaller = howManyLess(nums, value), both val.
Decide. If smaller < k, save result = value and move up with left = value + 1. Otherwise move down with right = value - 1. Return result at the end.
Finish kSmallest. Store kSmallestNum(nums, k) in kthNum, create result = IntArray(k) and keep the next free position in var i.
Copy the smaller numbers. Loop over nums and write each value less than kthNum to result[i++].
Pad with kthNum. While i < k, write kthNum to the remaining positions. This is needed because kthNum may appear several times, so fewer than k numbers can be strictly smaller. Return result.

The solution in Kotlin

class Solution {
    fun kSmallestNum(nums: IntArray, k: Int): Int {
        var left = nums[0]
        var right = nums[0]
        for (num in nums) {
            left = minOf(left, num)
            right = maxOf(right, num)
        }
        var result = left
        while (left <= right) {
            val value = left + (right - left) / 2
            val smaller = howManyLess(nums, value)
            if (smaller < k) {
                result = value
                left = value + 1
            } else {
                right = value - 1
            }
        }
        return result
    }
    fun howManyLess(nums: IntArray, value: Int): Int {
        var counter = 0
        for (num in nums) {
            if (value > num) {
                counter++
            }
        }
        return counter
    }
    fun kSmallest(nums: IntArray, k: Int): IntArray {
        val kthNum = kSmallestNum(nums, k)
        val result = IntArray(k)
        var i = 0
        for (value in nums) {
            if (value < kthNum) {
                result[i++] = value
            }
        }
        while (i < k) {
            result[i++] = kthNum
        }
        return result
    }
}

The Kotlin details worth a closer look:

IntArray(k). It creates an array of k zeros. Every position is overwritten before the function returns.
result[i++] = value. The post-increment uses the current i as the index and then moves i forward, so the writes land in consecutive positions.
Two variables named result. One is in kSmallestNum and one is in kSmallest. Each is local to its own function, so they don't clash.
minOf and maxOf. Standard-library functions for two values.
Helpers declared after use. kSmallestNum calls howManyLess although it comes later in the class. Member order doesn't matter in Kotlin.
Shorter count. nums.count { it < value } gives the same result as howManyLess in one expression.

Complexity

Time: O(n log R), where R is the distance between the largest and the smallest value. Each search iteration costs one pass over nums (O(n)), and halving the value range takes about log2(R) iterations. For an Int, R is below 2^32, so there are at most about 32 iterations, which is why this fits the lesson's O(n log n) budget for any realistic input. The final collection pass is O(n).

Space: O(k), for the returned array. Apart from that, the search uses O(1).

Edge cases and common mistakes

Duplicates of the k-th number. For [4, 4, 4, 4] and k = 2, nothing is strictly smaller than 4, so the result is padded: [4, 4]. Without padding the array would stay filled with zeros.
k equal to the length. For [9, 1] and k = 2 the k-th smallest number is 9, the numbers below it give [1], and the padding adds 9: [1, 9].
Counting with <= instead of <. The test smaller < k relies on strict counting. Counting equal values as smaller shifts the answer.
Collecting values less than or equal to kthNum. That can put more than k items into an array of size k and throw ArrayIndexOutOfBoundsException.
Searching over indices. The array is not sorted, so cutting the index range in half says nothing about the values.
Empty input. nums[0] fails on an empty array. The problem has at least k elements.

Testing

nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives [2, 1, 2, 3], which has the same values as [1, 2, 2, 3].
nums = [4, 4, 4, 4], k = 2 gives [4, 4].
nums = [9, 1], k = 2 gives [1, 9].

Keep going with KotlinAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free