Dutch National Flag Algorithm: Sort Colors in One Pass

"Sort an array that contains only 0s, 1s and 2s, in place, in one pass." It sounds like a sorting question, and a general sort would pass the tests. But with only three distinct values, sorting is really partitioning, and a partition can be done in one sweep with three pointers. The algorithm is known as the Dutch national flag problem, after the flag's three horizontal bands, and the detail that trips most people up is a single line that does not increment a counter.

The problem it solves

Given nums containing only 0, 1 and 2, rearrange it so all 0s come first, then 1s, then 2s. Use constant extra space and read each element once.

There are three obvious answers and each misses something:

• Call sort. Correct, but O(n log n) for a problem with three keys.
• Count, then rewrite. Count the 0s, 1s and 2s, then overwrite the array. O(n) and constant space, but two passes, and it only works because the values carry no other data. Sort records by a three-valued key and there is nothing to "rewrite" with.
• Two partitions. Move 0s to the front, then 2s to the back. Correct, also two passes.

The one-pass version moves the actual elements, so it works for records as well as for plain numbers.

The intuition

Keep the array split into four regions at all times:

• nums[0 .. low-1] — all 0, finished.
• nums[low .. i-1] — all 1, finished.
• nums[i .. high] — not inspected yet.
• nums[high+1 .. end] — all 2, finished.

At the start everything is unseen: low = i = 0, high = n - 1. Each step inspects nums[i] and shrinks the unseen region by one:

• A 0. Swap it with nums[low], the first 1 (or i itself, if no 1 has been seen yet). Both low and i advance. What came back to position i is a value already inspected, so it is safe to move past.
• A 1. It is already in the right region. Advance i.
• A 2. Swap it with nums[high] and shrink high. Do not advance i: the value that just arrived came from the unseen region, and nobody has looked at it yet.

The loop ends when i passes high and the unseen region is empty. That asymmetry, advancing after a swap with the front but not after a swap with the back, is the whole algorithm.

Watch it run

The animation sorts the lesson's [2, 0, 1, 2, 0]. The first value is a 2, so it is swapped to the back and high moves left; the i marker stays exactly where it was while the row changes under it. The 0 that arrived is then swapped into the front band. Watch the two bands under the row: the front band of 0s and the back band of 2s grow toward each other until the unseen middle disappears.

The code

The lesson's function, with the region invariant written into the comments:

def sort_colors(nums):
    low, i, high = 0, 0, len(nums) - 1
    while i <= high:                        # nums[i..high] is still unseen
        if nums[i] == 0:
            nums[low], nums[i] = nums[i], nums[low]
            low += 1
            i += 1                          # what came back was already seen
        elif nums[i] == 2:
            nums[high], nums[i] = nums[i], nums[high]
            high -= 1                       # i stays: the new value is unseen
        else:
            i += 1
    return nums

print(sort_colors([2, 0, 1, 2, 0]))         # [0, 0, 1, 2, 2]
print(sort_colors([2, 2, 2]))               # [2, 2, 2]
print(sort_colors([]))                      # []

The same loop partitions any array around a pivot value: smaller, equal, larger. Only the comparisons change. This is the partition step of three-way quicksort, which stops re-sorting runs of equal keys:

def three_way_partition(nums, pivot):
    low, i, high = 0, 0, len(nums) - 1
    while i <= high:
        if nums[i] < pivot:
            nums[low], nums[i] = nums[i], nums[low]
            low += 1
            i += 1
        elif nums[i] > pivot:
            nums[high], nums[i] = nums[i], nums[high]
            high -= 1
        else:
            i += 1
    return low, high                        # nums[low..high] all equal pivot

nums = [5, 1, 5, 9, 3, 5, 7]
print(three_way_partition(nums, 5), nums)   # (2, 4) [1, 3, 5, 5, 5, 7, 9]

The bug from the next section, reproduced. Advance i after the swap with high and a 0 that arrives from the back is never inspected:

def sort_colors_buggy(nums):
    low, i, high = 0, 0, len(nums) - 1
    while i <= high:
        if nums[i] == 0:
            nums[low], nums[i] = nums[i], nums[low]
            low += 1
        elif nums[i] == 2:
            nums[high], nums[i] = nums[i], nums[high]
            high -= 1
        i += 1                              # always advances: wrong
    return nums

print(sort_colors_buggy([2, 0, 1, 2, 0]))   # [0, 0, 1, 2, 2]  (lucky)
print(sort_colors_buggy([2, 1, 0]))         # [0, 1, 2]  (lucky again)
print(sort_colors_buggy([1, 2, 0]))         # [1, 0, 2]  (wrong)

Both correct functions against Python's sorted on 5,000 random arrays, and the partition checked against its own promise — everything before low smaller, everything after high larger, every value kept:

import random

random.seed(11)
ok = True
buggy_wrong = 0
for _ in range(5000):
    n = random.randint(0, 12)
    colors = [random.randint(0, 2) for _ in range(n)]
    ok &= sort_colors(colors[:]) == sorted(colors)
    buggy_wrong += sort_colors_buggy(colors[:]) != sorted(colors)

    vals = [random.randint(0, 9) for _ in range(n)]
    pivot = random.randint(0, 9)
    out = vals[:]
    lo, hi = three_way_partition(out, pivot)
    ok &= sorted(out) == sorted(vals)
    ok &= all(v < pivot for v in out[:lo])
    ok &= all(v == pivot for v in out[lo:hi + 1])
    ok &= all(v > pivot for v in out[hi + 1:])
print(ok, buggy_wrong > 0)                  # True True

The complexity

Every iteration either advances i or decreases high, so the unseen region shrinks by exactly one each time. That gives exactly n iterations, O(n) time, and at most one swap per iteration. Extra space is three integers, O(1).

The algorithm is not stable: two equal records can swap relative order, because a 2 swapped to the back jumps over everything between. If the elements are records and their original order matters, use a stable method such as counting sort with an output array, and pay O(n) extra space.

Where it goes wrong

• Advancing i after swapping with high. The value that arrives has never been inspected. If it is a 0, it is left stranded in the middle, as [1, 2, 0] shows above.
• Looping while i < high. The element at high is still unseen. A strict comparison skips it, so [1, 0] would come back unchanged.
• Not advancing i after swapping with low. It does not break correctness, because the value that returns is a 1 or the same 0, but it wastes iterations and hides that you understand why the front is safe.
• Hard-coding 0, 1, 2 when the question says "pivot". Interviewers often follow up with "partition around a value". Write the comparison version and the colours version is a special case.

How to say it in an interview

"There are only three values, so this is a partition, not a sort. I keep three pointers: everything before low is 0, everything after high is 2, and i to high is unseen. For each nums at i: a 0 swaps to low and both advance; a 1 just advances i; a 2 swaps with high and high shrinks, but i stays, because the value I pulled in from the back hasn't been looked at. The unseen region shrinks by one per step, so it's one pass, O(n) time, O(1) space. It isn't stable, and the same loop with a pivot is three-way quicksort's partition."

The same "swap it home, don't advance until the arrival is checked" rule drives cyclic sort, and the two-region version of the idea is move zeroes.