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

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}

Approach : Use the Dutch's flag algorithm given here. 

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35

#include <iostream> #include <cmath> using namespace std; void segregate(int *arr, int n){ int temp, low = 0, mid = 0, high = n - 1; while(mid <= high){ if(arr[mid] == 0){ temp = arr[low]; arr[low] = arr[mid]; arr[mid] = temp; mid++; low++; } else if(arr[mid] == 2){ temp = arr[high]; arr[high] = arr[mid]; arr[mid] = temp; mid; high--; } else { mid++; } } } int main(){ int arr[] = {0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1}; int n = sizeof(arr)/sizeof(arr[0]); segregate(arr, n); for(int i = 0; i < n; i++){ cout << arr[i] << " "; } cout << endl; return 0; }