Dutch national flag partition in Python

Why this matters

Many programs need to group items by a small label. A list of test results might need the failed ones first. A list of tasks might need the urgent ones at the front. Sorting the whole list works, but it does more than the job needs. A sort compares values with each other, and it can take many passes over the data.

This lesson solves a smaller version of that problem. Every value is one of three labels, written here as 0, 1 and 2. The goal is to move all the 0s to the front, all the 1s to the middle, and all the 2s to the end. The program walks through the list once, and it does not compare values with each other to decide where they go.

The problem has a name. It is called the Dutch national flag partition, after the three bands of colour on the Dutch flag. The Dutch computer scientist Edsger Dijkstra is often credited with posing it. A partition means splitting a collection into groups, with every item in exactly one group. The method you will write is a standard way to partition three groups in one pass.

This lesson uses the list basics from the one-dimensional lists lesson, and it swaps (exchanges) two items with a helper function. The out-of-bounds lab in this module covers how to read an index error, which is useful when you write the loop in this lesson.

The problem in plain words

Here is a small example. The input is a list of six numbers.

before:  2 0 2 1 1 0
after:   0 0 1 1 2 2

The output has the same six numbers, in the same counts, in a new order. The two 0s come first, then the two 1s, then the two 2s. The order inside each group does not matter, because all the 0s are equal to each other.

Two rules make the problem hard. First, you must not use a second list to build the answer. The lesson's program changes the list in place, which means it rearranges the same boxes. Second, you should walk through the list once. The built-in sorted function would also give the right answer, but it is a different tool. You will compare the two at the end.

The idea: three regions and three pointers

A pointer here means an index variable. It is not a special Python feature. It is a number that says which box you are looking at.

The program keeps three pointers, which are named low, mid and high. At any moment, the list is split into four parts. The table below describes each part.

At the start, the zeros and ones parts are empty, and the twos part is empty. The whole list is unknown. Each step looks at the box at mid and moves one of the four parts. When mid passes high, the unknown part is empty, and the list is sorted.

The rest of the lesson depends on this table. The three pointers are only numbers. Their meaning comes from the table. low is where the next 0 should go. mid is the box you are looking at now. high is where the next 2 should go, counting from the end.

Three rules, one for each value

Each step reads the value at mid and does one of three things.

• If the value is 0, swap it with the box at low. The 0 now sits in the zeros part, so move low up by one. Also move mid up by one. The box that came from low is a 1 that the program already passed, so it is safe to move past it.
• If the value is 1, leave it alone and move mid up by one. A 1 belongs in the middle, and the ones part grows to include it.
• If the value is 2, swap it with the box at high. The 2 now sits in the twos part, so move high down by one. Do not move mid. The value that came from high is not yet looked at, so mid must look at it next.

The third rule is the one students get wrong most often. The next section shows why it matters.

Worked example 1: the first input

Start with the list [2, 0, 2, 1, 1, 0]. The table below follows the program one step at a time. Each row shows the value at mid, what the program did, the three pointers after the step, and the list after the step.

At step 6, mid is 4 and high is 3, so mid is past high. The loop stops. The list is 0 0 1 1 2 2.

Look at step 1 closely. The value 2 at index 0 was swapped with the value 0 at index 5. The 0 moved to the front. The 2 moved to the back, and high moved down so the program would never touch that box again. The box at index 0 now holds a 0, and mid stays at 0 so the new value at index 0 is checked. It is a 0, so it is handled by rule 1 on the next step. If mid had moved up after a 2, the value that came from high would be skipped, and it might be a 2 that should have moved.

Worked example 2: a short input with a 1 in front

The second input is [1, 2, 0]. It is shorter, but it shows a case in which a 1 is passed over first.

The loop stops because mid (2) is past high (1). The 2 at the end was placed by rule 3 in step 2. The 0 was placed by rule 1 in step 3. The 1 ended in the middle, moving one box right in step 3 when the 0 was swapped in front of it.

Worked example 3: several swaps in a row

The third input is [2, 2, 0, 1, 0]. It starts with two 2s, so the first two steps both move the value to the back.

The final list is 0 0 1 2 2. Notice that step 2 handles a 0 that came from the back of the list. Every step shrinks the unknown part by one box.

The full program

This program puts the three rules into a function named sort_colours. The word "colours" is used because the labels are the colours of the flag. The function changes the list it is given. It returns nothing.

Output:

Before: [2, 0, 2, 1, 1, 0]
After:  [0, 0, 1, 1, 2, 2]

Three details in the code matter. The loop test is mid <= high, with a less-than-or-equal sign. When mid equals high, the box at mid is still unknown and must be checked. The swap function is separate so that each rule reads as one line. Python can also swap two items with the single line a[i], a[j] = a[j], a[i], and the function keeps that line in one place. The print calls print the list directly, so no helper is needed for printing.

Checking the result against a known answer

A correct function should match a result that you trust. In Python, sorted is that trusted result. It returns a new sorted list and leaves the original alone. The check below makes many random lists, sorts a copy of each with sorted, sorts the original with sort_colours, and compares the two.

Output:

cases=26000 failures=0

The sorted call makes a new list, so the original is sorted by sort_colours and the copy is the expected answer. The random generator is seeded with 42, so the same lists are made every time the check runs. The check covers lists of length 0 up to length 12, which includes the empty list and some short lists where the pointers meet at once. The seed is a number that starts the random sequence, so the same seed gives the same lists on every run. A check like this does not prove the function is correct for every possible input. It does show that the function agrees with the trusted sort on 26,000 lists. When you write a new algorithm, a check like this one is a cheap way to catch mistakes in the rules.

Common mistakes

The mistakes below are all easy to make. Each one shows what goes wrong.

The first mistake is to start high at the length of the list, not at the last index. The program then reads a box that does not exist. Save this version as main.py and run it:

def swap(a, i, j):
    a[i], a[j] = a[j], a[i]
 
 
def sort_colours(nums):
    low = 0
    mid = 0
    high = len(nums)
    while mid <= high:
        if nums[mid] == 0:
            swap(nums, low, mid)
            low += 1
            mid += 1
        elif nums[mid] == 1:
            mid += 1
        else:
            swap(nums, mid, high)
            high -= 1
 
 
nums = [2, 0, 2, 1, 1, 0]
print("Before:", nums)
sort_colours(nums)
print("After: ", nums)

The output is:

Before: [2, 0, 2, 1, 1, 0]
Traceback (most recent call last):
  File "main.py", line 23, in <module>
    sort_colours(nums)
    ~~~~~~~~~~~~^^^^^^
  File "main.py", line 17, in sort_colours
    swap(nums, mid, high)
    ~~~~^^^^^^^^^^^^^^^^^
  File "main.py", line 2, in swap
    a[i], a[j] = a[j], a[i]
                 ~^^^
IndexError: list index out of range

The first call to swap uses high, which is 6, so a[j] reads box 6. The last frame shows the failing line. The frames above it show how the program got there, from the main code to sort_colours to swap. The out-of-bounds lab in this module explains how to find this kind of bug from the traceback.

The second mistake is to use mid < high in the loop test. The loop then stops one step too soon, and the last unknown box is never checked. With the loop test changed in that one place, and the input changed to [1, 2, 0], the output is this.

[1, 0, 2]

No error is raised. The result is wrong and the program looks as if it worked. This is the most dangerous kind of bug, because nothing tells you it happened. A check against sorted, like the one above, would find it at once.

The third mistake is to move mid up after a swap with high. The value that came from the back of the list is then skipped. If that value is a 0, it stays in the twos part, which is wrong. The rule is that mid stays where it is after a 2.

The fourth mistake is to forget that the function does not check its input. The else branch treats every value that is not 0 or 1 as a 2. A value such as 5 is moved to the back, and the program does not report it. With the input changed to [5, 0], the output is this.

[0, 5]

The list is not sorted by the rule, because 5 is not one of the three labels, and the program gave no warning. If the input can contain other values, check each value before the loop and raise an error for anything outside 0, 1 and 2. For example, if any(v not in (0, 1, 2) for v in nums): raise ValueError("unexpected label") does the check in one line.

Practice

Try each exercise before you read the answer. Save each program as main.py and run it to check your answer.

Exercise 1

Trace the program by hand on the input [0, 1, 2]. What is the final list, and what is the value of high when the loop stops?

Show the answer

The final list is [0, 1, 2], and high is 1 when the loop stops. On the first step mid is 0 and the value is 0, so the program swaps box 0 with itself, and then moves low and mid to 1. On the second step the value is 1, so mid moves to 2. On the third step the value is 2, so the program swaps box 2 with itself and moves high down to 1. Now mid (2) is past high (1), and the loop stops.

Exercise 2

Why does the loop test use mid <= high, and not mid < high?

Show the answer

When mid equals high, the box at that index has not been looked at yet. It is still part of the unknown region, so it must be checked. The < version stops before that box is checked. The mistake section above shows the result on the input [1, 2, 0]: the output is [1, 0, 2], which is not sorted.

Exercise 3

In the full program, change the input to [2, 2]. Predict the After line before you run it.

Show the answer

The After line is After: [2, 2]. The list holds only 2s, so no value moves. The first step swaps box 0 with box 1, which leaves both values as 2, and high moves to 0. The second step swaps box 0 with itself and moves high to -1. Then mid (0) is past high (-1), and the loop stops. The list is sorted because every value is the same.

Exercise 4

Write a program that counts how many 0s, 1s and 2s are in [2, 0, 2, 1, 1, 0], using a single list of length 3 named counts. Print the three counts and then the list.

Show the answer

This program uses the value itself as the index. The value 0 goes to counts[0], the value 1 goes to counts[1], and the value 2 goes to counts[2].

Output:

zeros=2 ones=2 twos=2
[2, 2, 2]

The counts list has the same length as the number of labels, so the value can be used directly as the index. This is a counting list. It works only because the values are known to be 0, 1 or 2. Input from outside the program needs the check described in the common mistakes section.

Exercise 5

Explain in two or three sentences why this method reads the list only once, while sorted does more work. Then say when you would still use sorted.

Show the answer

Every step of sort_colours looks at one box at mid and either moves it into a finished part or passes it. Each step moves mid up or high down, so the loop runs at most as many times as there are boxes. sorted compares values with each other to find their order, and it works for any kind of value, not only three labels. Use sorted when the values are not limited to a few labels, or when you do not want to write the partition logic yourself.

Summary

• The Dutch national flag partition groups a list of 0s, 1s and 2s in one pass, with no second list.
• Three pointers, low, mid and high, split the list into finished zeros, finished ones, unknown values and finished twos.
• A 0 is swapped to the front and both low and mid move up. A 1 only moves mid up. A 2 is swapped to the back and only high moves down, so the swapped-in value is checked next.
• The loop test is mid <= high, because the box at mid is still unknown when the two pointers are equal.
• Check a new sorting method against sorted on many random inputs, including the empty list and very short lists.
• The function does not check its input, so values outside 0, 1 and 2 are moved without any error.