Partition Problems: 0s, 1s and 2s

A family of interview problems asks you to rearrange an array whose values come from a very small set. Separate 0s from 1s. Push all zeroes to the end. Sort an array of 0s, 1s and 2s.

Sorting solves all of them in O(n log n). All of them can also be done in O(n) with a single pass and no extra memory, because with only a handful of distinct values you do not need to compare anything. You can decide where a value belongs just by looking at it.

Segregating 0s and 1s

Two pointers walk inward. The left pointer advances while it sees 0s, the right advances backward while it sees 1s, and when both stop you have a 1 on the left and a 0 on the right, so swap them.

A counting approach also works: count the zeroes, then write that many 0s followed by 1s. It is equally O(n) and arguably simpler, but it destroys any other data attached to the elements, so it only works when the values really are just 0 and 1.

Moving zeroes to the end

The catch here is that the non-zero values must keep their original relative order, so the two-pointer swap-from-both-ends trick is wrong: it would scramble them.

Instead use a read pointer and a write pointer. Scan with the read pointer; every time you find a non-zero value, place it at the write pointer and advance the write pointer. Zeroes are simply skipped, and fall to the back as a consequence.

Switch the widget above to "Move zeroes" and watch the two pointers separate: the gap between them is exactly the number of zeroes seen so far.

The Dutch national flag problem

Three values, one pass. This is Dijkstra's Dutch national flag problem, named for the flag's three horizontal bands, and it is the one genuinely tricky member of the family.

Keep three pointers. Everything before low is 0, everything after high is 2, and mid scans the unknown middle:

That last row is where almost everyone gets it wrong. The value swapped in from high has never been examined, so mid must stay put and look at it on the next iteration. Advancing past it can leave a 0 or a 2 stranded in the middle region.

The loop ends when mid passes high, not when it reaches the end of the array: everything above high is already known to be 2.

Why this matters beyond the puzzle

Three-way partitioning is exactly what quick sort needs when the input has many duplicate keys. A standard two-way partition puts values equal to the pivot on one side, so an array of all-equal elements produces maximally unbalanced splits and degrades to O(n²). Partitioning into less-than, equal-to and greater-than instead makes that same input O(n), because the entire equal region is finished in one step.

In code

// Dutch national flag: 0s, 1s and 2s in a single pass.
function sortColors(a) {
  let low = 0, mid = 0, high = a.length - 1;

  while (mid <= high) {
    if (a[mid] === 0) {
      [a[low], a[mid]] = [a[mid], a[low]];
      low++; mid++;
    } else if (a[mid] === 2) {
      [a[mid], a[high]] = [a[high], a[mid]];
      high--;
      // mid does NOT advance: the value swapped in is still unexamined.
    } else {
      mid++;
    }
  }
  return a;
}

// Move zeroes to the end, keeping the other values in order.
function moveZeroes(a) {
  let write = 0;
  for (let read = 0; read < a.length; read++) {
    if (a[read] !== 0) {
      [a[write], a[read]] = [a[read], a[write]];
      write++;
    }
  }
  return a;
}

Check yourself