Sort Colors

medium · arrays, two-pointers, partition

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 nums in-place.

Example

nums = [2,0,2,1,1,0]
After sort: [0,0,1,1,2,2]

Constraints

  • 1 <= len(nums) <= 100_000
  • nums[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)

  1. Initialize low = 0, mid = 0, high = len(nums)-1.
  2. While mid <= high:
    • If nums[mid] == 0, swap with nums[low], increment both low and mid.
    • If nums[mid] == 1, just increment mid.
    • If nums[mid] == 2, swap with nums[high], decrement high (do not increment mid).
  3. 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)

Trace 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 mid after swapping with high (the new value at mid is 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.

Return to Two Pointers · Module quiz

Run tests to see results
No issues detected
    Join Discord