Sort Colors

Difficulty: Medium

Problem

A bag holds n marbles, stored in the array nums as 0 (red), 1 (white) or 2 (blue). Rearrange nums in place so all the reds come first, then the whites, then the blues. The library sort is off limits, and a single pass over the array is the stretch goal.

Example

Input: nums = [1, 2, 0, 1, 0, 2, 2]
Output: [0, 0, 1, 1, 2, 2, 2]
Explanation: Two 0s (red) go first, then two 1s (white), then three 2s (blue).

Brute-force approach

Count occurrences of each color, then overwrite the array.

Steps

  1. Count occurrences of 0s, 1s, and 2s
  2. Overwrite array: first count0 zeros
  3. Then count1 ones
  4. Then count2 twos

Time complexity: O(n) · Space complexity: O(1)

Trade-offs

  • Requires two passes through the array
  • Not a true one-pass algorithm
  • Can we do better with three pointers?

Optimal approach

Key insight: Maintain three regions: [0, low) contains 0s, [low, mid) contains 1s, (high, n-1] contains 2s. The unknown region is [mid, high].

Dutch National Flag: use three pointers to partition in one pass.

Steps

  1. Initialize low=0, mid=0, high=n-1
  2. While mid <= high:
  3. If nums[mid]=0: swap with low, increment both
  4. If nums[mid]=1: just increment mid
  5. If nums[mid]=2: swap with high, decrement high only

Time complexity: O(n) · Space complexity: O(1)

Pseudocode

function sortColors(nums):
    low = 0       // boundary for 0s
    mid = 0       // current element
    high = n - 1  // boundary for 2s
    
    while mid <= high:
        if nums[mid] == 0:
            swap(nums[low], nums[mid])
            low++, mid++
        else if nums[mid] == 1:
            mid++
        else:  // nums[mid] == 2
            swap(nums[mid], nums[high])
            high--

Hints

  • Use three pointers: low, mid, high
  • 0s go to the left, 2s go to the right, 1s stay in middle
  • Dutch National Flag algorithm

Complexity questions

What is the time complexity of the counting approach?

O(n). One pass to count + one pass to fill = O(2n) = O(n).

What is the time complexity of the Dutch National Flag?

O(n). Single pass through the array = O(n).

What is the space complexity of the Dutch National Flag?

O(1). Only 3 pointer variables, regardless of input size.

Why don't we advance mid after swapping with high?

The swapped element from high is unknown and needs checking. The element swapped from high hasn't been checked yet - it could be 0, 1, or 2.

Check your understanding of Sort Colors

What do the numbers 0, 1, 2 represent?

Colors: red, white, blue. 0=red, 1=white, 2=blue. We sort by color.

What is the expected order after sorting?

0s, 1s, 2s (ascending). All 0s first, then 1s, then 2s.

For input [2, 0, 2, 1, 1, 0], what is the expected output?

[0, 0, 1, 1, 2, 2]. All 0s first, then 1s, then 2s, sorted in ascending order.

Related problems

Keep practising