Sort Colors: One-Pass Three-Way Partition of 0s, 1s and 2s

Problem

You get a list containing only the values 0, 1 and 2 (think of three colours). Sort it in place so all 0s come first, then 1s, then 2s, without calling a library sort. Ideally use one pass and constant extra space. This is widely known as LeetCode 75, Sort Colors.

Examples

[2, 0, 1, 2, 1, 0]   ->  [0, 0, 1, 1, 2, 2]
[1, 1]               ->  [1, 1]
[2, 0]               ->  [0, 2]

Approach 1: brute force (counting sort)

Count each value, then overwrite the list.

def sort_colors_count(nums):
    counts = [0, 0, 0]
    for v in nums:
        counts[v] += 1
    i = 0
    for value in range(3):
        for _ in range(counts[value]):
            nums[i] = value
            i += 1

Complexity: O(n) time over two passes, O(1) extra space. This is a perfectly good answer; the follow-up asks for a single pass.

Approach 2: optimal (Dutch national flag)

Key insight: keep the array in four regions: [0, low) holds 0s, [low, mid) holds 1s, [mid, high] is unexamined, and (high, end] holds 2s. Each step shrinks the unexamined region by one.

Walkthrough on [2, 0, 1, 2, 1, 0] (low = 0, mid = 0, high = 5):

mid value	Action	Array after	low, mid, high
2	swap mid and high	[0, 0, 1, 2, 1, 2]	0, 0, 4
0	swap mid and low	[0, 0, 1, 2, 1, 2]	1, 1, 4
0	swap mid and low	[0, 0, 1, 2, 1, 2]	2, 2, 4
1	advance	[0, 0, 1, 2, 1, 2]	2, 3, 4
2	swap mid and high	[0, 0, 1, 1, 2, 2]	2, 3, 3
1	advance	[0, 0, 1, 1, 2, 2]	2, 4, 3: stop

def sort_colors(nums):
    low, mid, high = 0, 0, len(nums) - 1
    while mid <= high:
        if nums[mid] == 0:
            nums[low], nums[mid] = nums[mid], nums[low]
            low += 1
            mid += 1
        elif nums[mid] == 1:
            mid += 1
        else:
            nums[mid], nums[high] = nums[high], nums[mid]
            high -= 1

Complexity: O(n) time in one pass, O(1) extra space.

Tests

import random

def run(f, arr):
    arr = list(arr)
    f(arr)
    return arr

for f in (sort_colors, sort_colors_count):
    assert run(f, [2, 0, 1, 2, 1, 0]) == [0, 0, 1, 1, 2, 2]
    assert run(f, [1, 1]) == [1, 1]                  # one colour
    assert run(f, [2, 0]) == [0, 2]
    assert run(f, []) == []                          # empty
    assert run(f, [2]) == [2]                        # single
    assert run(f, [2, 2, 2, 0, 0, 0]) == [0, 0, 0, 2, 2, 2]   # no 1s
    assert run(f, [0, 1, 2]) == [0, 1, 2]            # already sorted

big = [2, 1, 0] * 30_000
assert run(sort_colors, big) == sorted(big)          # large input

random.seed(25)
for _ in range(400):
    arr = [random.randint(0, 2) for _ in range(random.randint(0, 12))]
    assert run(sort_colors, arr) == run(sort_colors_count, arr) == sorted(arr)

Edge cases and pitfalls

• After swapping a 2 to high, do not advance mid: the value that came back from high has not been examined.
• After swapping a 0 to low, it is safe to advance mid, because the value that came back from low is a 1 (or the same cell).
• The loop condition is mid <= high, not mid < high; the cell at high is still unexamined.
• Negative values or values above 2 are out of contract; ask before assuming the input is clean.

Where this shows up in data engineering

Three-way partitioning is how some sort routines handle heavy duplicate keys efficiently: values equal to the pivot are grouped once instead of being re-sorted. Splitting a batch into accepted, quarantined and rejected rows in a single pass is the same idea applied to records.