LESSON 262 OF 347 · DIVIDE AND CONQUER

Smallest K Integers in C#: O(n) Quickselect with a Random Pivot

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

C# · DIVIDE AND CONQUER

Smallest K in C#, explained

Problem overview

The method KSmallest receives an int[] named nums and an integer k. It must return the k smallest values as an int[], in any order. For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the answer is the values 1, 2, 2 and 3. Duplicates count separately, which is why 2 appears twice. The arrays [1, 2, 2, 3] and [2, 1, 3, 2] are both correct.

The lesson asks for O(n) time and O(log n) extra space. Sorting the whole array would take O(n log n), which is more than needed: the answer doesn't have to be ordered, so most of that work is wasted.

Approach: partition, then keep only the side that matters

This is the Quickselect idea, a cousin of Quick Sort. Partition the array around a random pivot so that every value less than or equal to the pivot comes before it and every larger value comes after it. The pivot lands at its final sorted position, pivot. That tells us exactly how many values are at most the pivot: pivot - left + 1 of them, counting the pivot.

If that count is at least k, all k smallest values are inside [left, pivot]. The right side is irrelevant, so recurse on [left, pivot - 1].
If the count is less than k, everything in [left, pivot] belongs to the answer. Only the remaining k - (pivot - left + 1) values are still missing, and they are the smallest of [pivot + 1, right].

Quick Sort recurses into both sides. Quickselect recurses into only one, so the work shrinks geometrically: about n + n/2 + n/4 + ..., which is O(n) on average with a random pivot. When the search ends, the first k positions of the array hold the k smallest values, in no particular order.

Take [3, 1, 2] with k = 2 and suppose the pivot is 3. Everything is at most 3, so the pivot ends at index 2 and the count is 3, which is at least 2. The search continues in [0, 1], where the remaining work puts 1 and 2 first. The first two positions then hold 1 and 2.

Step by step

Declare SmallestKNums. It takes the array, the inclusive bounds left and right of the range it works on, and k. It returns nothing and rearranges nums in place.
Return early on tiny ranges. If left >= right, there is at most one value and nothing to do.
Declare Partition above it, returning the pivot's final index as an int.
Pick a random pivot. randomIndex is left plus a random offset from 0 to right - left. A random pivot keeps the expected time linear even when the input is already sorted.
Declare Swap, a small helper that exchanges two elements through a temporary variable.
Move the pivot to the end of the range with Swap, so the scan can partition everything else first.
Record the pivot value and the boundary. value is the pivot, and correctIndex starts at left: the position where the next value that is not larger than the pivot belongs.
Scan from left to right, both included. The pivot is visited last.
Swap small values into place. If 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.
Partition in SmallestKNums. After the early return, store the result of Partition in pivot.
Count at least k? Recurse on [left, pivot - 1] with the same k. The pivot is the largest value of the range up to pivot, so it can stay where it is.
Otherwise, recurse on [pivot + 1, right] with k - (pivot - left + 1).
In KSmallest, call SmallestKNums on the whole array, from 0 to nums.Length - 1.
Return the first k values as a new array.

The solution in C#

class Solution {
    public void Swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    public int Partition(int[] nums, int left, int right) {
        int randomIndex = left + (int)(new Random().NextDouble() * (right - left + 1));
        Swap(nums, randomIndex, right);

        int value = nums[right];
        int correctIndex = left;
        for (int i = left; i <= right; ++i) {
            if (nums[i] <= value) {
                Swap(nums, i, correctIndex++);
            }
        }
        return correctIndex - 1;
    }

    public void SmallestKNums(int[] nums, int left, int right, int k) {
        if (left >= right) {
            return;
        }
        int 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));
        }
    }

    public int[] KSmallest(int[] nums, int k) {
        SmallestKNums(nums, 0, nums.Length - 1, k);
        return nums.Take(k).ToArray();
    }
}

The C# details worth a closer look:

correctIndex++ as an argument. Post-increment passes the current value and then adds one. Swap(nums, i, correctIndex++) therefore does the same as swapping with correctIndex and incrementing on the next line.
LINQ's Take and ToArray. Take(k) is lazy: it describes the first k elements without copying them. ToArray() runs it and builds the int[] the method returns. These extension methods come from System.Linq. If k is larger than the array, Take simply stops at the end instead of throwing.
The cast in the random index. NextDouble() returns a value from 0 up to but not including 1. Scaling by the range size and casting with (int) truncates the fraction. Random.Shared.Next(left, right + 1) would give the same kind of index, with an exclusive upper bound.
The input array is changed. nums is partly rearranged, because arrays are passed by reference. Only the returned copy holds the answer.

Complexity

Time: O(n) expected. The first partition costs about n steps. With a random pivot, the range that remains is, on average, a constant fraction of the previous one, so the costs add up to a geometric series that stays proportional to n. Copying the final k values adds O(k). Unlucky pivots can make a single run slower, and the absolute worst case is O(n²), but randomness makes that very unlikely.

Space: O(log n) expected for the recursion stack, plus the O(k) result array that the method has to return anyway.

Edge cases and common mistakes

k equal to the array length. Every value belongs to the answer, so however the search rearranges the array, Take(k) returns all of its values.
k = 0. The count is always at least 0, so the search shrinks to the left each time and Take(0) returns an empty array.
Duplicates. The partition test is <=, so equal values stay on the small side. The answer for the example contains the value 2 twice, as it should.
Forgetting to subtract in the else branch. After keeping pivot - left + 1 values, only k minus that many are still missing. Passing the original k would ask for too many.
Expecting sorted output. Only the first k positions are guaranteed to hold the right values. Sort the returned array if you need them in order.
Recursing into both sides. That turns the algorithm back into Quick Sort and loses the linear running time.

Testing

nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 returns 1, 2, 2 and 3 in some order.
nums = [3, 1, 2], k = 0 returns an empty array.
nums = [3, 1, 2], k = 3 returns all three values.

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

Start learning free