Sort 0s, 1s and 2s

Sort 0s, 1s and 2s (also known as Sort Colors or the Dutch National Flag problem) asks you to sort an array containing only three distinct values — in place, in a single pass. It is a classic test of the three-pointer partition trick.

The slow way first

The easy answer is counting sort: count how many 0s, 1s, and 2s there are, then overwrite the array with that many of each. It is O(n) time and O(1) space — but it takes two passes (one to count, one to write).

A general comparison sort like a.sort() is even worse: O(n log n), and it ignores the gift we were handed. We only have three possible values, so we should be able to place every element correctly in a single sweep.

The idea: partition into three regions

Keep three pointers that carve the array into four zones:

• everything left of low is a settled 0,
• everything right of high is a settled 2,
• the band between mid and high is unknown,
• and mid is the cursor scanning that unknown band.

For each value at mid, do one of three things and the regions grow until the unknown band disappears:

The subtle part: when you swap a 2 down to high, mid does not advance — the value that came back from high is still unknown, so you must look at it again.

Walk through it

Step through the animation. Watch low build the wall of 0s from the left and high build the wall of 2s from the right while mid sweeps the middle. Settled boxes turn green. The loop ends the instant mid passes high — there is nothing unknown left to examine.

Pseudocode

low, mid = 0, 0
high = last index
while mid <= high:
    if a[mid] == 0:
        swap a[low] and a[mid]    # send the 0 left
        low += 1; mid += 1
    else if a[mid] == 1:
        mid += 1                  # 1 is already in the middle
    else:                         # a[mid] == 2
        swap a[mid] and a[high]   # send the 2 right
        high -= 1                 # do NOT move mid

The Python solution

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

• low, mid, high start the unknown band as the whole array.
• The loop runs while mid <= high — once they cross, the unknown band is empty.
• a[mid] == 0: swap it down to the low wall, then advance both low and mid (the value that came up to mid was already scanned, so it is safe to step over).
• a[mid] == 1: it already belongs in the middle band, so only mid moves.
• a[mid] == 2: swap it up to the high wall and shrink high. mid stays put because the swapped-in value is unexamined.

Complexity

Each step moves either mid forward or high backward, and they only ever close toward each other, so the loop runs at most n times. No extra arrays — the sort happens inside the input.

When this pattern shows up

Practice