Given an array A[] consisting 0s, 1s and 2s, write a function that sorts A[]. The functions should put all 0s first, then all 1s and all 2s in last.

Example
Input = {0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1};
Output = {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2}

The problem is similar to our old post Segregate 0s and 1s in an array, and both of these problems are variation of famous Dutch national flag problem.

The problem was posed with three colours, here `0′, `1′ and `2′. The array is divided into four sections:

  1. a[1..Lo-1] zeroes (red)
  2. a[Lo..Mid-] ones (white)
  3. a[Mid..Hi] unknown
  4. a[Hi+1..N] twos (blue)

The unknown region is shrunk while maintaining these conditions

  1. Lo := 1; Mid := 1; Hi := N;
  2. while Mid <= Hi do
    1. Invariant: a[1..Lo-1]=0 and a[Lo..Mid-1]=1 and a[Hi+1..N]=2; a[Mid..Hi] are unknown.
    2. case a[Mid] in
      • 0: swap a[Lo] and a[Mid]; Lo++; Mid++
      • 1: Mid++
      • 2: swap a[Mid] and a[Hi]; Hi–