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

Software Development
1

Hello Everyone,

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

Input : {0, 1, 2, 0, 1, 2} Output : {0, 0, 1, 1, 2, 2} 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 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-1] ones (white)
    3. a[Mid…Hi] unknown
    4. a[Hi+1…N] twos (blue)
    5. If the ith element is 0 then swap the element to the low range, thus shrinking the unknown range.
    6. Similarly, if the element is 1 then keep it as it is but shrink the unknown range.
    7. If the element is 2 then swap it with an element in high range.

  • Algorithm:

    1. Keep three indices low = 1, mid = 1 and high = N and there are four ranges, 1 to low (the range containing 0), low to mid (the range containing 1), mid to high (the range containing unknown elements) and high to N (the range containing 2).
    2. Traverse the array from start to end and mid is less than high. (Loop counter is i)
    3. If the element is 0 then swap the element with the element at index low and update low = low + 1 and mid = mid + 1
    4. If the element is 1 then update mid = mid + 1
    5. If the element is 2 then swap the element with the element at index high and update high = high – 1 and update i = i – 1. As the swapped element is not processed
    6. Print the output array.

  • Dry Run:
    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].

  • Implementation:

// C++ program to sort an array

// with 0, 1 and 2 in a single pass

#include <bits/stdc++.h>

using namespace std;

// Function to sort the input array,

// the array is assumed

// to have values in {0, 1, 2}

void sort012( int a[], int arr_size)

{

int lo = 0;

int hi = arr_size - 1;

int mid = 0;

// Iterate till all the elements

// are sorted

while (mid <= hi) {

switch (a[mid]) {

// If the element is 0

case 0:

swap(a[lo++], a[mid++]);

break ;

// If the element is 1 .

case 1:

mid++;

break ;

// If the element is 2

case 2:

swap(a[mid], a[hi--]);

break ;

}

}

}

// Function to print array arr[]

void printArray( int arr[], int arr_size)

{

// Iterate and print every element

for ( int i = 0; i < arr_size; i++)

cout << arr[i] << " " ;

}

// Driver Code

int main()

{

int arr[] = { 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 };

int n = sizeof (arr) / sizeof (arr[0]);

sort012(arr, n);

cout << "array after segregation " ;

printArray(arr, n);

return 0;

}

  • Output:

array after segregation 0 0 0 0 0 1 1 1 1 1 2 2

  • Complexity Analysis:
  • Time Complexity: O(n).
    Only one traversal of the array is needed.
  • Space Complexity: O(1).
    No extra space is required.