vizlearn

Sort an array of 0s, 1s and 2s

Interview QuestionsBy Updated

Three pointers carve the array into four regions: settled 0s, settled 1s, unexamined, and settled 2s. A 0 swaps down, a 2 swaps up, a 1 stays. One pass, O(1) space — and the one subtlety is that after swapping a 2 the middle pointer must not advance.

Overview

The counting alternative

Count the 0s, 1s and 2s, then overwrite the array. Two passes, trivially correct, and usually the first answer. The question asks for one pass to rule it out — and because counting sort does not generalise to sorting objects by a three-way key, which the partition does.

Lists & arraysCoding problemMedium

Step through it

What to watch

  • Four regions: settled 0s, settled 1s, unexamined, settled 2s.
  • After swapping a 2 down, mid stays — the incoming value is unseen.
  • The loop ends when mid passes hi, not the array end.

Say this out loud

"Dutch national flag. Three pointers - low, mid, high. 0 swaps to the low region and both advance, 2 swaps to the high region and only high moves, 1 just advances mid. One pass, O(1) space."

Sort an array of 0s, 1s and 2s

Sort an array containing only 0, 1 and 2 in a single pass, in place.

The invariant

Everything before lo is 0. Everything from lo to mid is 1. Everything after hi is 2. Between mid and hi is unexamined. The loop restores that invariant on every step and stops when the unexamined region is empty.

Being able to state the invariant is most of the answer here. Writing the three branches from memory without it is how the mid bug appears.

The one asymmetry

After swapping a 0 into the low region, the value that came back is from the region already known to be 1s, so mid can safely advance. After swapping a 2 into the high region, the value that came back is from the unexamined region — so mid must stay and look at it.

Advancing in both cases is the classic bug: the array comes out almost sorted, with stray 2s left in the middle. It is exactly what the question is testing.

Run it in Python

The three-way partition with its invariant asserted on every iteration, the buggy version that advances mid in both branches, and both checked against sorted() on random input.

sort_colors.pyPython 3
Output
How the code works
  1. mid += 1 after a 0 swapSafe, because the value swapped back comes from the region already known to hold 1s. Nothing unexamined arrives at mid.
  2. no mid += 1 after a 2 swapThe value swapped back comes from the unexamined region, so it has to be looked at. Advancing here is the bug the second function demonstrates on real input.
  3. while mid <= hiThe loop ends when the unexamined region is empty, not at the end of the array — everything past hi is already settled.
  4. the three assert linesThe invariant, checked rather than described. Stating it is most of the answer to this question; writing the branches without it is how the bug appears.
Change one thing
  • Delete the asserts and add a fourth colour. The approach does not extend — three-way partitioning is specifically three-way.
  • Sort objects by a three-way key instead of raw integers. The partition still works; counting does not.
Where this runs
Real CPython, compiled to WebAssembly and running on your own machine — nothing is uploaded. The first run takes a few seconds while the interpreter downloads; after that it is immediate. Need more room, or want to paste your own attempt? Use the Python compiler.


    



Check yourself
0 of 3
Answer without scrolling back up.
  1. After swapping a 2 from mid to the high region, why must mid stay?
  2. The regions maintained by the invariant are:
  3. Why is the counting approach not the accepted answer?





 Cheat sheet
Sort an array of 0s, 1s and 2s
Three pointers carve the array into four regions: settled 0s, settled 1s, unexamined, and settled 2s. A 0 swaps down, a 2 swaps up, a 1 stays. One pass, O(1) space — and the one subtlety is that after swapping a 2 the middle pointer must not advance.
INTERVIEW · vizlearn.in/interview/sort-colors-dutch-national-flag.html
More in Interview
Interview3SumInterviewSearch in a rotated sorted arrayInterviewSliding window maximumInterviewlist vs tuple vs deque vs array — which and why?
About the author
Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.
About VizLearnGitHubLinkedInContact