Algorithms‎ > ‎Array‎ > ‎

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

June 30, 2010

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–

— Dutch National Flag Algorithm, or 3-way Partitioning —

Part way through the process, some red, white and blue elements are known and are in the “right” place. The section of unknown elements, a[Mid..Hi], is shrunk by examining a[Mid]:

 `|0 0 0 1 1 1 ? ? ? ? 2 2 2^     ^     ^|     |     |Lo    Mid   Hi`

Examine a[Mid]. There are three possibilities: a[Mid] is (0) red, (1) white or (2) blue.
Case (0) a[Mid] is red, swap a[Lo] and a[Mid]; Lo++; Mid++

 `0 0 0 0 1 1 1 ? ? ? 2 2 2^     ^   ^|     |   |Lo    Mid Hi`

Case (1) a[Mid] is white, Mid++

 `0 0 0 1 1 1 1 ? ? ? 2 2 2^       ^   ^|       |   |Lo      Mid Hi`

Case (2) a[Mid] is blue, swap a[Mid] and a[Hi]; Hi–

 `0 0 0 1 1 1 ? ? ? 2 2 2 2^     ^   ^|     |   |Lo    Mid Hi`

Continue until Mid>Hi.

 `#include` `/* Function to swap *a and *b */``void` `swap(``int` `*a, ``int` `*b);` `void` `sort012(``int` `a[], ``int` `arr_size)``{``   ``int` `lo = 0;``   ``int` `hi = arr_size - 1;``   ``int` `mid = 0;` `   ``while``(mid <= hi)``   ``{``      ``switch``(a[mid])``      ``{``         ``case` `0:``           ``swap(&a[lo++], &a[mid++]);``           ``break``;``         ``case` `1:``           ``mid++;``           ``break``;``         ``case` `2:``           ``swap(&a[mid], &a[hi--]);``           ``break``;``      ``}``   ``}``}` `/* UTILITY FUNCTIONS */``void` `swap(``int` `*a, ``int` `*b)``{``  ``int` `temp = *a;``  ``*a = *b;``  ``*b = temp;``}` `/* Utility function to print array arr[] */``void` `printArray(``int` `arr[], ``int` `arr_size)``{``  ``int` `i;``  ``for` `(i = 0; i < arr_size; i++)``    ``printf``(``"%d "``, arr[i]);``  ``printf``(``"\n"``);``}` `/* driver program to test */``int` `main()``{``  ``int` `arr[] = {0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1};``  ``int` `arr_size = ``sizeof``(arr)/``sizeof``(arr);``  ``int` `i;` `  ``sort012(arr, arr_size);` `  ``printf``(``"array after segregation "``);``  ``printArray(arr, arr_size);` `  ``getchar``();``  ``return` `0;``}`

Time Complexity: O(n)

The above code performs unnecessary swaps for inputs like 0 0 0 0 1 1 1 2 2 2 2 2 : lo=4 and mid=7 and hi=11. In present code: first 7 exchanged with 11 and hi become 10 and mid is still pointing to 7. again same operation is till the mid <= hi. But it is really not required. We can put following loop before switch condition to make sure that hi is pointing to location which is not 2 so that it would eliminate unnecessary swaps of 2.

 `while` `((a[hi]==2) && hi >= mid)``    ``–hi;``if` `(hi < mid)``    ``break``;`

Thanks to rka for suggesting this change.