Sort Colors [Medium] — Dutch National Flag Algorithm

Problem Statement

Given an array with only 0, 1, and 2, sort it in-place so all 0s come first, then all 1s, then all 2s. One pass, O(1) extra space.

Input: [2,0,2,1,1,0] → Output: [0,0,1,1,2,2]

Intuition

Three pointers: lo (boundary of 0s), mid (current element), hi (boundary of 2s). Sweep mid left to right, swapping as needed.

Solutions

C

void sortColors(int* nums, int n) {
    int lo = 0, mid = 0, hi = n - 1;
    while (mid <= hi) {
        if (nums[mid] == 0)      { int t=nums[lo]; nums[lo++]=nums[mid]; nums[mid++]=t; }
        else if (nums[mid] == 2) { int t=nums[hi]; nums[hi--]=nums[mid]; nums[mid]=t; }
        else                     { mid++; }
    }
}

C++

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

Java

public void sortColors(int[] nums) {
    int lo = 0, mid = 0, hi = nums.length - 1;
    while (mid <= hi) {
        if (nums[mid] == 0)      { int t=nums[lo]; nums[lo++]=nums[mid]; nums[mid++]=t; }
        else if (nums[mid] == 2) { int t=nums[hi]; nums[hi--]=nums[mid]; nums[mid]=t; }
        else                     { mid++; }
    }
}

JavaScript

var sortColors = function(nums) {
    let lo = 0, mid = 0, hi = nums.length - 1;
    while (mid <= hi) {
        if (nums[mid] === 0)      { [nums[lo++], nums[mid++]] = [nums[mid], nums[lo]]; }
        else if (nums[mid] === 2) { [nums[mid], nums[hi--]] = [nums[hi], nums[mid]]; }
        else                      { mid++; }
    }
};

Python

def sortColors(nums: list[int]) -> None:
    lo = mid = 0
    hi = len(nums) - 1
    while mid <= hi:
        if nums[mid] == 0:
            nums[lo], nums[mid] = nums[mid], nums[lo]
            lo += 1; mid += 1
        elif nums[mid] == 2:
            nums[mid], nums[hi] = nums[hi], nums[mid]
            hi -= 1
        else:
            mid += 1

Complexity

Time: O(n) — single pass
Space: O(1)