Color Segregation Challenge

You're in a technical phone screen at Apple when the interviewer says, "Given an array of 0s, 1s, and 2s, sort it in-place without using a library sort function." You immediately recognize this as the classic Sort Colors problem, a staple at companies like Adobe, Google, Microsoft, and Amazon. On Firecode, we call it the "Color Segregation Challenge," and the optimal solution uses a technique that dates back to one of computer science's greatest minds: Edsger Dijkstra's Dutch National Flag algorithm.

TL;DR

Sort an array containing only 0s, 1s, and 2s in-place using three pointers: red (boundary of 0s), white (current scanner), and blue (boundary of 2s). Scan from left to right with white. When you see a 0, swap it to the front and advance both red and white. When you see a 1, just advance white. When you see a 2, swap it to the back and pull blue inward, but do not advance white because the swapped-in element still needs inspection. This sorts the array in O(n) time and O(1) space with a single pass.

Why This Problem Matters

Sort Colors tests a few skills at once: whether you understand in-place algorithms, whether you can manage multiple pointers without losing track of invariants, and whether you know when to advance a pointer versus when to hold still. The counting sort two-pass approach is straightforward, but interviewers expect the single-pass Dutch National Flag solution. It shows up in variations across partitioning problems, three-way quicksort, and any scenario where you need to segregate data into three buckets without extra memory.

Understanding the Problem

Given an array nums of n elements where each element is 0, 1, or 2 (representing red, white, and blue), reorder them in-place so all reds come first, then whites, then blues.

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

Constraints:

• 1 <= n <= 300
• Each element is 0, 1, or 2
• Must sort in-place without using a library sort
• Target: O(n) time, O(1) space

Edge Cases

1. Single element: Already sorted, return as-is
2. All same value: [1, 1, 1] - no swaps needed
3. Already sorted: [0, 1, 2] - algorithm handles it with no unnecessary work
4. Reverse sorted: [2, 1, 0] - requires the most swaps

The Naive Approach: Two-Pass Counting

The simplest approach is counting sort: count how many 0s, 1s, and 2s exist, then overwrite the array.

// Two-pass counting approach - works but not optimal
int count0 = 0, count1 = 0, count2 = 0;
for (int num : nums) {
    if (num == 0) count0++;
    else if (num == 1) count1++;
    else count2++;
}
int i = 0;
while (count0-- > 0) nums[i++] = 0;
while (count1-- > 0) nums[i++] = 1;
while (count2-- > 0) nums[i++] = 2;

This runs in O(n) time and O(1) space, but it makes two passes. Interviewers will ask: "Can you do it in one pass?" That is where the Dutch National Flag algorithm comes in.

Solution Approach: The Dutch National Flag Algorithm

Dijkstra designed this algorithm for partitioning an array into three groups. The idea is to maintain three pointers that divide the array into four regions:

• Everything before red: confirmed 0s
• Between red and white: confirmed 1s
• Between white and blue: unprocessed elements
• Everything after blue: confirmed 2s

The white pointer scans through the unprocessed region. At each step, you look at nums[white] and decide what to do:

1. If nums[white] == 0: This red element belongs at the front. Swap it with nums[red], then advance both red and white. The element that lands at white after the swap was previously between red and white, so it must be a 1 (already processed). That is why advancing white is safe.

2. If nums[white] == 1: This white element is already in the correct region. Just advance white.

3. If nums[white] == 2: This blue element belongs at the back. Swap it with nums[blue], then pull blue inward. Do NOT advance white, because the element swapped in from position blue is unprocessed and needs to be checked.

The loop terminates when white > blue, meaning the unprocessed region is empty and the array is sorted.

Why Not Advance White After a Blue Swap?

This is the part that trips up most candidates. When you swap nums[white] with nums[blue], you pulled an unknown element from the right side of the array into position white. That element could be a 0, 1, or 2. If you advanced white past it, you would leave it in the wrong place. After a red swap, the element at white came from between red and white, which means it was already classified as a 1. So advancing is safe in that case.

Tracing Through [2, 0, 2, 1, 1, 0]

Starting state:

Steps 0-2: The first element is 2, so swap with position blue (index 5). Then position 0 holds 0 (was swapped from the end), so swap with red and advance both.

Steps 3-5: Another 0 at position 1, swap with red and advance. Position 2 holds 2, swap with blue.

Steps 6-7: Two 1s in a row - just advance white each time until white > blue.

Here is the complete flow at a glance:

Implementation

Prefer a different language? Jump to solutions in other languages.

public class Solution {
  public int[] sortColors(int[] nums) {
    int red = 0, white = 0, blue = nums.length - 1;

    while (white <= blue) {
      if (nums[white] == 0) {
        // Swap current element to the red region
        int temp = nums[red];
        nums[red] = nums[white];
        nums[white] = temp;
        red++;
        white++;
      } else if (nums[white] == 1) {
        // Already in the correct region
        white++;
      } else {
        // Swap current element to the blue region
        int temp = nums[white];
        nums[white] = nums[blue];
        nums[blue] = temp;
        blue--;
        // Do NOT increment white - the swapped element needs inspection
      }
    }

    return nums;
  }
}

The three pointers start at their natural positions: red and white at the beginning (index 0), and blue at the end (index nums.length - 1). The while loop continues until white passes blue, at which point every element has been placed in its correct region.

Step-by-Step Execution

For input [2, 0, 1]:

Result: [0, 1, 2].

Complexity Analysis

Time Complexity: O(n)

Each iteration of the loop either advances white forward or pulls blue backward (or both in the red case). The gap between white and blue shrinks by at least one every iteration, so the loop runs at most n times. Each swap is O(1), giving O(n) total.

Space Complexity: O(1)

Three integer variables (red, white, blue) and one temporary variable for swapping. No extra arrays, hash maps, or data structures. The sorting is purely in-place.

Comparison with Alternatives

The Dutch National Flag algorithm wins on all fronts for this specific problem.

Common Pitfalls

1. Advancing white after a blue swap: The element swapped in from blue has not been examined. Skipping it can leave 0s or 2s in the wrong region.

2. Wrong loop condition: Using white < blue instead of white <= blue misses the case where white and blue point to the same element. That element still needs processing.

3. Off-by-one on blue initialization: blue should start at nums.length - 1, not nums.length. Starting at nums.length would cause an out-of-bounds access on the first blue swap.

4. Forgetting to advance red AND white after a red swap: If you advance only red, the white pointer falls behind and the algorithm loops forever on the same element.

Interview Tips

When you get this problem in an interview:

1. Mention both approaches: Start by acknowledging the counting sort approach works in O(n) with two passes. Then say you can do it in one pass using the Dutch National Flag technique. This shows you can think of multiple solutions before diving in.

2. Draw the four regions: Sketch the red/white/unprocessed/blue regions on the whiteboard. Walk through why each swap preserves the invariants. Interviewers love seeing candidates reason about correctness.

3. Explain the white pointer asymmetry: Proactively explain why you advance white after a red swap but not after a blue swap. This is the part where most candidates get tripped up, and bringing it up shows you genuinely understand the algorithm.

4. Test with edge cases: Run through [0], [2, 2, 2], and [2, 1, 0] to demonstrate your solution handles boundaries correctly.

5. Know the history: If asked where this algorithm comes from, mention Dijkstra and the Dutch flag. It is a nice touch that shows depth of knowledge.

Key Takeaways

• The Dutch National Flag algorithm sorts an array of three distinct values in O(n) time and O(1) space using a single pass with three pointers.
• The critical invariant: elements before red are 0s, between red and white are 1s, between white and blue are unprocessed, and after blue are 2s.
• Never advance the white pointer after swapping with blue, because the swapped-in element has not been inspected yet. After a red swap, the landed element is guaranteed to be 1, so advancing is safe.
• The counting sort two-pass approach is a valid O(n) solution, but interviewers expect the single-pass variant to demonstrate understanding of pointer invariants.
• This three-way partitioning pattern is the same concept behind three-way quicksort and appears whenever you need to bucket data into three groups in-place.

Practice Makes Perfect

Once you have Sort Colors down, try these related problems to reinforce the pattern:

• Move Zeroes (two-way partition variant)
• Sort an array of 0s and 1s (simpler two-pointer version)
• Three-way quicksort partition
• Wiggle Sort (alternating order using similar swap logic)

This problem and thousands of others are available on Firecode, where over 50,000 users have prepared for technical interviews and landed roles at top tech companies. Whether you are just getting started or aiming for a senior role at a FAANG company, building fluency with in-place partitioning algorithms like this one will serve you well.

Happy coding, and may all your pointers land in the right region!

Solutions in Other Languages

Python

class Solution:
    def sort_colors(self, nums):
        red, white, blue = 0, 0, len(nums) - 1

        while white <= blue:
            if nums[white] == 0:
                nums[red], nums[white] = nums[white], nums[red]
                red += 1
                white += 1
            elif nums[white] == 1:
                white += 1
            else:
                nums[white], nums[blue] = nums[blue], nums[white]
                blue -= 1

        return nums

JavaScript

sortColors(nums) {
  let red = 0, white = 0, blue = nums.length - 1;

  while (white <= blue) {
    if (nums[white] === 0) {
      [nums[red], nums[white]] = [nums[white], nums[red]];
      red++;
      white++;
    } else if (nums[white] === 1) {
      white++;
    } else {
      [nums[white], nums[blue]] = [nums[blue], nums[white]];
      blue--;
    }
  }

  return nums;
}

TypeScript

sortColors(nums: number[]): number[] {
  let red = 0, white = 0, blue = nums.length - 1;

  while (white <= blue) {
    if (nums[white] === 0) {
      [nums[red], nums[white]] = [nums[white], nums[red]];
      red++;
      white++;
    } else if (nums[white] === 1) {
      white++;
    } else {
      [nums[white], nums[blue]] = [nums[blue], nums[white]];
      blue--;
    }
  }

  return nums;
}

C++

std::vector<int> sortColors(std::vector<int> &nums) {
  int red = 0, white = 0, blue = nums.size() - 1;

  while (white <= blue) {
    if (nums[white] == 0) {
      std::swap(nums[red], nums[white]);
      red++;
      white++;
    } else if (nums[white] == 1) {
      white++;
    } else {
      std::swap(nums[white], nums[blue]);
      blue--;
    }
  }

  return nums;
}

Go

Go uses low, mid, and high as the pointer names and a switch statement, which reads naturally for this kind of branching.

func (s *Solution) SortColors(nums []int) []int {
    low, mid, high := 0, 0, len(nums)-1

    for mid <= high {
        switch nums[mid] {
        case 0:
            nums[low], nums[mid] = nums[mid], nums[low]
            low++
            mid++
        case 1:
            mid++
        case 2:
            nums[mid], nums[high] = nums[high], nums[mid]
            high--
        }
    }

    return nums
}

Scala

Scala uses match for pattern matching on the current element.

def sortColors(nums: Array[Int]): Array[Int] = {
  var (red, white, blue) = (0, 0, nums.length - 1)

  while (white <= blue) {
    nums(white) match {
      case 0 =>
        val temp = nums(red)
        nums(red) = nums(white)
        nums(white) = temp
        red += 1
        white += 1
      case 1 =>
        white += 1
      case 2 =>
        val temp = nums(white)
        nums(white) = nums(blue)
        nums(blue) = temp
        blue -= 1
    }
  }

  nums
}

Swift

Swift provides swapAt for clean in-place swapping on arrays.

func sortColors(_ nums: [Int]) -> [Int] {
    var nums = nums
    var red = 0
    var white = 0
    var blue = nums.count - 1

    while white <= blue {
        if nums[white] == 0 {
            nums.swapAt(red, white)
            red += 1
            white += 1
        } else if nums[white] == 1 {
            white += 1
        } else {
            nums.swapAt(white, blue)
            blue -= 1
        }
    }

    return nums
}

Kotlin

Kotlin's also function provides a concise swap idiom.

fun sortColors(nums: IntArray): IntArray {
    var red = 0
    var white = 0
    var blue = nums.size - 1

    while (white <= blue) {
        if (nums[white] == 0) {
            nums[red] = nums[white].also { nums[white] = nums[red] }
            red++
            white++
        } else if (nums[white] == 1) {
            white++
        } else {
            nums[white] = nums[blue].also { nums[blue] = nums[white] }
            blue--
        }
    }

    return nums
}

Rust

Rust requires careful handling of signed/unsigned types since blue can go below zero conceptually.

pub fn sort_colors(&self, mut nums: Vec<i32>) -> Vec<i32> {
    let mut red = 0_usize;
    let mut white = 0_usize;
    let mut blue = nums.len() as i32 - 1;

    while (white as i32) <= blue {
        if nums[white] == 0 {
            nums.swap(white, red);
            red += 1;
            white += 1;
        } else if nums[white] == 1 {
            white += 1;
        } else {
            nums.swap(white, blue as usize);
            blue -= 1;
        }
    }

    nums
}

Ruby

Ruby's parallel assignment makes swaps especially concise.

def sort_colors(nums)
  red = 0
  white = 0
  blue = nums.length - 1

  while white <= blue
    if nums[white] == 0
      nums[red], nums[white] = nums[white], nums[red]
      red += 1
      white += 1
    elsif nums[white] == 1
      white += 1
    else
      nums[white], nums[blue] = nums[blue], nums[white]
      blue -= 1
    end
  end

  nums
end

Dart

List<int> sortColors(List<int> nums) {
  int red = 0, white = 0, blue = nums.length - 1;

  while (white <= blue) {
    if (nums[white] == 0) {
      int temp = nums[red];
      nums[red] = nums[white];
      nums[white] = temp;
      red++;
      white++;
    } else if (nums[white] == 1) {
      white++;
    } else {
      int temp = nums[white];
      nums[white] = nums[blue];
      nums[blue] = temp;
      blue--;
    }
  }

  return nums;
}

C#

public int[] sortColors(int[] nums) {
    int red = 0, white = 0, blue = nums.Length - 1;

    while (white <= blue) {
        if (nums[white] == 0) {
            int temp = nums[red];
            nums[red] = nums[white];
            nums[white] = temp;
            red++;
            white++;
        } else if (nums[white] == 1) {
            white++;
        } else {
            int temp = nums[white];
            nums[white] = nums[blue];
            nums[blue] = temp;
            blue--;
        }
    }

    return nums;
}

PHP

public function sortColors(array $nums): array {
    $red = 0;
    $white = 0;
    $blue = count($nums) - 1;

    while ($white <= $blue) {
        if ($nums[$white] === 0) {
            $temp = $nums[$red];
            $nums[$red] = $nums[$white];
            $nums[$white] = $temp;
            $red++;
            $white++;
        } elseif ($nums[$white] === 1) {
            $white++;
        } else {
            $temp = $nums[$white];
            $nums[$white] = $nums[$blue];
            $nums[$blue] = $temp;
            $blue--;
        }
    }

    return $nums;
}