Sort Colors

mediumTwo Pointers48 stepsJava · Python · C++

Sort an array of 0s, 1s and 2s in one pass. Keep three regions: 0s before lo, 2s after hi, and 1s between lo and mid. mid scans the unknown region and swaps each value into its region.

▶ Watch the animated trace

Traced input: nums = [2, 0, 2, 1, 1, 0] → [0, 0, 1, 1, 2, 2]. Original problem: neetcode.io/problems/sort-colors

Complexity

BestO(n)

AverageO(n)

WorstO(n)

SpaceO(1)

Each element is examined once, in place.

Key idea

Invariant: [0, lo) are 0s, [lo, mid) are 1s, (hi, n) are 2s, [mid, hi] is unknown.
A 0 swaps to lo; both lo and mid advance because the swapped-in value is a known 1.
A 2 swaps to hi; only hi moves, because the value swapped into mid hasn't been checked yet.
A 1 is already in place, so just advance mid.

Common mistakes

Advancing mid after swapping with hi skips an unchecked value.
Looping while (mid < hi) misses the last element; it must be mid <= hi.
Counting sort is two passes; the follow-up asks for one.

Solution

Java

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

void swap(int[] a, int i, int j) {
    int t = a[i]; a[i] = a[j]; a[j] = t;
}

Python

def sortColors(nums: List[int]) -> None:
    lo, mid, hi = 0, 0, 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] == 1:
            mid += 1
        else:
            nums[mid], nums[hi] = nums[hi], nums[mid]
            hi -= 1

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]);
            lo++;
            mid++;
        } else if (nums[mid] == 1) {
            mid++;
        } else {
            swap(nums[mid], nums[hi]);
            hi--;
        }
    }
}

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
Already sorted Swaps only with itself.	`nums = [0, 1, 2]`	`[0, 1, 2]`
Reverse sorted The 2 goes to the end first, then the 0 to the front.	`nums = [2, 1, 0]`	`[0, 1, 2]`
Single colour Only `mid` moves.	`nums = [1, 1, 1]`	`[1, 1, 1]`
Single element One swap with itself and `hi` drops below `mid`.	`nums = [2]`	`[2]`
Only 0s and 2s The empty middle region is fine.	`nums = [2, 0, 2, 0]`	`[0, 0, 2, 2]`