LESSON 263 OF 348 · DIVIDE AND CONQUER

Smallest K in TypeScript: O(n) Expected with 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

TYPESCRIPT · DIVIDE AND CONQUER

Smallest K in TypeScript, explained

Problem overview

The method kSmallest receives an array of positive integers, nums, and a count k. It returns the k smallest values, in any order. Duplicates count separately.

For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the sorted array starts 1, 2, 2, 3, so those are the four values to return. [2, 1, 3, 2] is just as correct. The statement asks for O(n) time and O(log n) extra space.

Approach: partition, then keep only the side that matters

Sorting the whole array and taking the first k values works, but it costs O(n log n). A heap of size k costs O(n log k). Neither reaches O(n). This lesson uses the idea behind Quick Sort, without the part that sorts everything.

Partitioning picks a pivot and rearranges a range so that every value not larger than the pivot sits before it and every larger value after it. The pivot then has its final sorted position, pivot. So the range [left, pivot] holds exactly 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 in [left, pivot]. Only [left, pivot - 1] still needs work, since the pivot is the largest value of the group.
If the count is smaller than k, the whole group belongs to the answer, and we need the k - (pivot - left + 1) smallest values from [pivot + 1, right].

Quick Sort recurses into both sides. This approach, usually called Quickselect, recurses into one. With a good pivot the range roughly halves each time, so the work is about n + n/2 + n/4 + ..., which adds up to O(n). A random pivot makes that the expected cost on any input order.

When the recursion ends, the first k positions of nums hold the k smallest values, unsorted among themselves, and the method returns a copy of them.

Step by step

Declare smallestKNums with nums, left, right and k. It returns void and rearranges nums in place so the k smallest values of the range come first.
Base case. When left >= right, the range has at most one value, so return.
Declare partition with nums, left and right, returning the pivot's final index.
Choose a random pivot with left + Math.floor(Math.random() * (right - left + 1)). Math.random() is in [0, 1), so this covers each index from left to right with equal chance.
Write a swap helper that exchanges two positions through a temporary variable.
Move the pivot to the end by swapping randomIndex and right.
Save the pivot value in value, and start correctIndex at left. It is the next slot for a value not larger than the pivot.
Scan i from left to right, both included, so the pivot is visited last.
Move small values left. When nums[i] <= value, swap it with nums[correctIndex] and advance correctIndex.
Return correctIndex - 1. The pivot was the last value moved, so it sits just before correctIndex.
Back in smallestKNums, store the result of partition in pivot.
If pivot - left + 1 >= k, recurse on [left, pivot - 1] with the same k.
Otherwise, recurse on [pivot + 1, right] with k - (pivot - left + 1).
In kSmallest, call smallestKNums on the whole array.
Return nums.slice(0, k), a new array of the first k values. The end index of slice is exclusive.

The solution in TypeScript

class Solution {
    swap(nums: number[], i: number, j: number): void {
        const temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    partition(nums: number[], left: number, right: number): number {
        const randomIndex = left + Math.floor(Math.random() * (right - left + 1));
        this.swap(nums, randomIndex, right);
        const value = nums[right];
        let correctIndex = left;
        for (let i = left; i <= right; ++i) {
            if (nums[i] <= value) {
                this.swap(nums, i, correctIndex++);
            }
        }
        return correctIndex - 1;
    }

    smallestKNums(nums: number[], left: number, right: number, k: number): void {
        if (left >= right) {
            return;
        }
        const pivot = this.partition(nums, left, right);
        if (pivot - left + 1 >= k) {
            this.smallestKNums(nums, left, pivot - 1, k);
        } else {
            this.smallestKNums(nums, pivot + 1, right, k - (pivot - left + 1));
        }
    }

    kSmallest(nums: number[], k: number): number[] {
        this.smallestKNums(nums, 0, nums.length - 1, k);
        return nums.slice(0, k);
    }
}

Here is one partition worked by hand on the example, assuming the random choice picks the last element, 5, as the pivot. After the scan the array is [5, 3, 2, 1, 3, 2, 5, 6, 7, 9] and partition returns 6. The range [0, 6] holds 7 values, which is at least k = 4, so the next call looks only at [0, 5]. Another partition there narrows it further, until the first four positions hold 1, 2, 2, 3 in some order.

TypeScript details worth knowing:

this.swap(nums, i, correctIndex++). The post-increment passes the current value of correctIndex and then adds one, so the swap uses the old index and the next one is ready afterwards.
const versus let. const stops a variable from being reassigned. It does not freeze an array: nums is a parameter, and nums[i] = ... still changes the caller's array.
The input is reordered. Arrays are objects, so the caller's nums ends up rearranged. The returned array is a separate copy made by slice.
Types are erased. number[] and void are checked when compiling and disappear from the JavaScript that runs.

Complexity

Let n be the length of nums.

Time: O(n) expected. One partition of a range of length m costs O(m). With a random pivot the range shrinks geometrically in expectation, so the total is a small constant times n. The worst case is O(n²) if the pivots are repeatedly extreme, which random choice makes very unlikely. Copying with slice adds O(k).

Space: O(log n) expected for the recursion stack, plus O(k) for the returned array. No other array is created.

Edge cases and common mistakes

k = 0. The result is []. The recursion keeps going left, ends at the base case, and slice(0, 0) is empty.
k larger than the length. slice stops at the end of the array, so the result holds every value. For [4, 2, 9] and k = 5 you get all three.
Duplicates. partition uses <=, so copies of the pivot value go to the left side. With [2, 2, 2, 2] and k = 3 the result is [2, 2, 2]. When many values are equal, though, the splits become uneven and the speed drops toward n².
Not reducing k on the right side. After moving to [pivot + 1, right], the values already in [left, pivot] are part of the answer, so k must shrink by their count.
Dropping the early return. Without left >= right, an empty or one-element range would recurse without end.
Expecting sorted output. The returned values are the right ones but in no particular order.
Using sort() without a comparator as a shortcut. nums.sort() compares as strings, so [10, 9, 1] becomes [1, 10, 9]. A numeric sort needs (a, b) => a - b.

Testing

Sort the result before comparing, since the order is free.

nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives [1, 2, 2, 3].
nums = [4, 2, 9], k = 0 gives [].
nums = [4, 2, 9], k = 5 gives [2, 4, 9].
nums = [7], k = 1 gives [7].
nums = [2, 2, 2, 2], k = 3 gives [2, 2, 2].

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

Start learning free