Sort Colors

Sort Colors is sorting with only three possible values, and it is the cleanest example of beating the O(n log n) bound. A counting sort already does it in O(n) with two passes. The interview question is the one-pass version, known as the Dutch national flag partition, after the flag's three bands.

The same three-pointer partition is the heart of three-way quicksort and quickselect, which Wiggle Sort II uses later in this section. Learn its invariant here, where the values are simple.

The problem

An array holds only the values 0, 1 and 2 (think red, white and blue). Rearrange it in place so all 0s come first, then all 1s, then all 2s. Do not use the library sort.

• [2, 0, 2, 1, 1, 0] → [0, 0, 1, 1, 2, 2].
• [2, 1, 0] → [0, 1, 2].

Constraints: 1 ≤ n ≤ 10⁵ (an empty array should also work). The follow-up interviewers ask for: one pass, O(1) extra space.

Clarifying questions

• In place, or may I return a new array? In place.
• Only 0, 1 and 2? Yes. (If more values were possible, this becomes a counting sort or a partition around a pivot.)
• Does stability matter? No — equal values are indistinguishable.
• One pass required? That is the follow-up; start with what works.

Approach 1: the simple way — count, then overwrite

The library sort works but is O(n log n) and is ruled out. The simple honest answer is a counting sort: count each colour, then write them back.

1def sort_colors_counting(nums: list[int]) -> None:2    """Two passes: count 0s, 1s and 2s, then overwrite."""3    counts = [0, 0, 0]4    for x in nums:5        counts[x] += 16    i = 07    for colour in range(3):8        for _ in range(counts[colour]):9            nums[i] = colour10            i += 1

This is O(n) time and O(1) space — already optimal in big-O. It is not too slow; it just does not meet the stated follow-up of a single pass. It also only works because the items are their values. If each item were a record with a colour field, overwriting would lose the records, while a partition moves them.

The key insight

Keep the array split into four regions at every moment, using three pointers:

At the start everything is unknown. Each step looks at nums[mid], the first unknown value, and moves it into its region:

• 0: swap it with nums[low], the first 1 (or mid itself if there are no 1s yet). Both low and mid move right. The value that came back from low was a 1 we had already seen, so it is safe to step past it.
• 1: it is already in place at the end of the middle region. Move mid right.
• 2: swap it with nums[high] and move high left. Do not move mid: the value that came back from high has never been looked at.

When mid passes high, the unknown region is empty and the array is sorted.

Approach 2: optimised — the Dutch national flag

1def sort_colors(nums: list[int]) -> None:2    """One pass, O(1) space: the Dutch national flag partition."""3    low, mid, high = 0, 0, len(nums) - 14    while mid <= high:5        if nums[mid] == 0:6            nums[low], nums[mid] = nums[mid], nums[low]7            low += 18            mid += 1                   # what came from low is a 1 we already saw9        elif nums[mid] == 1:10            mid += 111        else:12            nums[mid], nums[high] = nums[high], nums[mid]13            high -= 1                  # do NOT move mid: the new value is unseen

Dry run on [2, 0, 2, 1, 1, 0]:

At step 1 the 2 is swapped to the back and a 0 comes forward — unseen, so mid stays at 0 and step 2 examines it. At step 4 the swap brings a 1 from index 4 to index 2, and step 5 examines it. After step 6, mid (4) has passed high (3): done.

Complexity: O(n) time — every step either moves mid right or high left, so there are at most n steps. O(1) space — three integers.

Edge cases

• Empty array: high = -1, the loop never runs.
• One element: one step, nothing moves.
• All the same colour, [0, 0, 0]: each step swaps nums[mid] with itself and moves on.
• Only 2s then 0s, [2, 2, 0]: the first 2 swaps with the 0, which is then examined at the same mid; result [0, 2, 2].

Follow-ups

• k colours instead of 3? Counting sort in O(n + k), or repeat the partition around each value in turn.
• Items are records with a colour field? The same three-pointer code works; swap the records. The counting version does not.
• Partition around any pivot value? This is exactly partition3 from the core lesson: "less than pivot" plays the role of 0, "equal" of 1, "greater" of 2.