Sort Colors
Sort Colors
You are given an array containing only 0, 1, and 2. Sort it in-place so that all 0s come first, then 1s, then 2s.
You must do it in a single pass with O(1) extra space.
Why this is worth learning
Four regions, not three pointers alone, explain the Dutch flag algorithm: zeros, ones, unknown values, and twos.
Before coding: one-minute experiment
At state [1,2,0], low=0, mid=1, high=2, what must happen after swapping the 2 with the 0?
Function signature
func SortColors(nums []int)
Inputs / outputs
nums: slice containing only 0, 1, 2.- Return: nothing; modify
numsin-place.
Example
nums = [2,0,2,1,1,0]
After sort: [0,0,1,1,2,2]
Constraints
1 <= len(nums) <= 100_000nums[i]is 0, 1, or 2.
Core idea (near-solution)
Use the Dutch National Flag algorithm with three regions:
[0, low)are all 0s[low, mid)are all 1s(high, end]are all 2s
Invariant: elements between mid and high are unknown and need processing.
Algorithm (step-by-step)
- Initialize
low = 0,mid = 0,high = len(nums)-1. - While
mid <= high:- If
nums[mid] == 0, swap withnums[low], increment bothlowandmid. - If
nums[mid] == 1, just incrementmid. - If
nums[mid] == 2, swap withnums[high], decrementhigh(do not incrementmid).
- If
- Done when
mid > high.
Correctness sketch
Maintain nums[:low] all zero, nums[low:mid] all one, nums[mid:high+1] unknown, and nums[high+1:] all two. A zero swapped into low extends the zero region; the displaced value is a known one unless low==mid, so both low and mid can advance. A one extends the one region by advancing mid. A two swapped to high extends the two region, but the incoming value came from the unknown region, so mid must stay. Each branch reduces the unknown region by one. When it is empty, the three finalized regions form the sorted result without losing or adding values.
Trace (hand walkthrough)
Example input:
nums = [2,0,2,1,1,0]
After sort: [0,0,1,1,2,2]
Walk:
- low=0, mid=0, high=5; nums=[2,0,2,1,1,0]
- mid=0 val=2 -> swap with high: [0,0,2,1,1,2], high=4
- mid=0 val=0 -> swap with low: [0,0,2,1,1,2], low=1, mid=1
- mid=1 val=0 -> swap with low: [0,0,2,1,1,2], low=2, mid=2
- mid=2 val=2 -> swap with high: [0,0,1,1,2,2], high=3
- mid=2 val=1 -> mid=3; mid=3 val=1 -> mid=4 stop
Pitfalls to avoid
- Incrementing
midafter swapping withhigh(the new value atmidis unprocessed). - Using extra arrays (violates O(1) space).
Debugging lens
The potential function is high-mid+1, not the number of loop iterations that advance mid. Test [1,2,0] and long runs of 2s followed by 0s. Inputs outside {0,1,2} are outside this contract.
Complexity
- Time:
O(n) - Extra space:
O(1)
Implementation notes (Go)
- In-place swaps are done with
nums[i], nums[j] = nums[j], nums[i].
Reconstruction checkpoint
At state [1,2,0], low=0, mid=1, high=2, what must happen after swapping the 2 with the 0?
Model answer
The array becomes [1,0,2], high becomes 1, and mid stays 1. The incoming 0 still needs processing; incrementing mid would leave [1,0,2] unsorted.