Dutch national flag: sort 0, 1, 2

Sort an array that only holds three distinct values in one pass with three pointers.

Time O(n)Space O(1)Code C++, Python and C#

Watch it animatedStep through it, with the running line highlighted. Free with an account.

The idea

Keep three regions: zeros before lo, ones between lo and mid, twos after hi. mid inspects the unknown middle part and hands each value to its region.

Code

C++PythonC#

void sortColors(vector<int>& a) {
    int lo = 0, mid = 0, hi = (int)a.size() - 1;
    while (mid <= hi) {
        if (a[mid] == 0) { swap(a[lo], a[mid]); lo++; mid++; }
        else if (a[mid] == 1) mid++;
        else { swap(a[mid], a[hi]); hi--; }
    }
}

def sort_colors(a: list[int]) -> None:
    lo, mid, hi = 0, 0, len(a) - 1
    while mid <= hi:
        if a[mid] == 0: a[lo], a[mid] = a[mid], a[lo]; lo += 1; mid += 1
        elif a[mid] == 1: mid += 1
        else: a[mid], a[hi] = a[hi], a[mid]; hi -= 1
    # sorted in place

public static void SortColors(int[] a) {
    int lo = 0, mid = 0, hi = a.Length - 1;
    while (mid <= hi) {
        if (a[mid] == 0) { (a[lo], a[mid]) = (a[mid], a[lo]); lo++; mid++; }
        else if (a[mid] == 1) mid++;
        else { (a[mid], a[hi]) = (a[hi], a[mid]); hi--; }
    }
}

Tips

After swapping with hi, do not advance mid. The value that arrived from the right has not been inspected yet.
This three-way split is also how quick sort handles many duplicate keys.

In the .NET library

Array.Sort would work but costs O(n log n). This is the O(n) answer interviewers ask for.