Partition Problems: 0s, 1s and 2s

A whole family of coding-interview questions is built on one premise: you are handed an array whose entries are drawn from a tiny set of possible values, and you have to shuffle them into contiguous groups. One version asks you to separate the 0s from the 1s, another wants every zero pushed to the back while the rest keep their places, and a third hands you a mixture of 0s, 1s and 2s to put in order.

Sorting any of these arrays into shape is an O(n log n) operation, and it always works, but it does far more than the problem requires. Because the range of possible values is so narrow, you never have to compare one element against another; the value you are holding already tells you which group it belongs to. A single sweep from one end of the array to the other is therefore enough to finish the job, and you can carry it out in place without borrowing a scratch array.

Segregating 0s and 1s

Imagine one pointer starting at the front of the array and another at the back, each walking toward the centre. The front pointer keeps advancing for as long as it lands on 0s, since those are already on the side where they belong, and the back pointer mirrors that behaviour, sliding left past every 1 it meets. Sooner or later both of them stall: the front pointer is sitting on a 1 that ought to be near the back, and the back pointer is sitting on a 0 that ought to be near the front. Exchanging that pair corrects both positions in one move, after which the two pointers resume their walk until they cross.

Step the run above one frame at a time. The two highlighted bars are the pointers; the caption under them names the swap or the skip that each frame performs, and the comparison count never moves because nothing here is ever compared.

A cruder approach reaches the same result just as quickly. Walk the array once to count how many zeroes it contains, then rewrite it from the beginning as a run of that many 0s followed by 1s for the remainder. It is almost impossible to get wrong, but it carries a real cost: it discards everything each element was carrying beyond its bare value, so it is only safe when the entries really are nothing more than the numbers 0 and 1. The moment those numbers are keys attached to larger records, overwriting them throws away data you cannot recover.

Moving zeroes to the end

This version looks almost identical, but it adds one demanding rule: everything that is not a zero has to end up in the same order it started in. That single constraint rules out the swap-from-both-ends idea, which moves elements across long distances and would leave the surviving values scrambled relative to their original sequence.

The method that preserves the order uses two pointers travelling in the same direction rather than converging. The first, which you can think of as the reader, visits every position in turn. The second, the writer, holds the place where the next value worth keeping should be dropped. Whenever the reader finds something other than a zero, that value is copied back to the writer's slot and the writer moves forward by one. Zeroes are passed over without comment, and because every non-zero value gets packed toward the front in its original order, the zeroes are left to accumulate at the end of the array as a side effect.

Step through the run above and watch the writer trail behind the reader. The gap that opens between the two pointers is never accidental: at any moment it is exactly the count of zeroes the reader has already passed over, which is also how many positions every surviving value has shifted toward the front.

The Dutch national flag problem

Introduce a third value and the exercise becomes distinctly harder. Arranging an array of 0s, 1s and 2s in a single sweep is the classic problem known as Dijkstra's Dutch national flag problem, a name taken from the three horizontal bands of the Netherlands' flag, and it is the one member of this family where the bookkeeping is genuinely easy to get subtly wrong.

The method tracks three indices, low, mid, and high, which between them divide the array into four regions. Everything to the left of low has already been settled as 0, everything to the right of high has been settled as 2, and the stretch running from low up to mid holds the 1s placed so far. That leaves the section between mid and high as the only part still unexamined, and mid is the marker that advances through it one element at a time. The action you take at each step is decided entirely by the value currently under mid:

That last row is the one people stumble on. When the swap happens against the high end, the element pulled into the slot under mid has come straight from unexplored ground and could just as easily be a 0, a 1 or a 2, so mid has to hold its ground and judge the newcomer on the following pass. If mid strides past it instead, a stray 0 or 2 can be left marooned in the middle band, and the final arrangement is quietly wrong.

The scan is over the instant mid climbs past high, which normally happens well before mid would have reached the physical end of the array. Carrying on would be pointless, because every position beyond high was already fixed as a 2 while the pointers were getting there.

Why this matters beyond the puzzle

Quick sort runs into this exact issue whenever its input is dense with repeated keys. The standard two-way partition drops every element equal to the pivot onto a single side, so an array in which every value is the same splits as unevenly as it possibly can and the running time collapses to O(n²). Splitting into three parts instead, one below the pivot, one equal to it and one above, clears the entire block of equal elements in one move, and that same pathological input then finishes in linear time.

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