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

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}

A simple solution is discussed in this(Sort an array of 0s, 1s and 2s (Simple Counting)) post.

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-1] 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 —

filter_none

edit
close

play_arrow

link
brightness_4
code

<script type="text/javascript"
// <![CDATA[
function dnf(N)
    var a = new Array(N + 1);
    var i;
    document.theForm3.opt.value = '';
    document.theForm3.trace.value = '';
      
    for(i = 1; i <= N; i++)
        a[i] = Math.round(100 * Math.random()) % 3;
    for(i = 1; i <= N; i++) 
        document.theForm3.opt.value += ' ' + a[i];
      
    var Lo = 1, Mid = 1, Hi = N;
    while( Mid - 0 <= Hi )
    {
        // inv: a[1..Lo-1]=0 && a[Lo..Mid-1]=1 && 
        //      a[Hi+1..N]=2, a[Mid..Hi] unknown.
        if(a[Mid] - 0 == 0)
        
            var temp = a[Mid]; 
                a[Mid] = a[Lo]; 
                a[Lo] = temp;
            Lo++; // the zeros' section grew
            Mid++;
        }
        else if(a[Mid]-0 == 1)
            Mid++; // ones' section grew
        else// ==2
        
            var temp = a[Hi]; 
                a[Hi] = a[Mid]; 
                a[Mid] = temp;
            Hi--; // the twos' section grew
        }
      
        for(i = 1; i <= N; i++)
            if( i - 0 < Mid || i - 0 > Hi)
            document.theForm3.trace.value += ' ' + a[i];
            else
            document.theForm3.trace.value += ' ';
        document.theForm3.trace.value += 'n';
    }
}
// ]]></script>                    

chevron_right


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]:

DNF1

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++

DNF2


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

DNF3


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

DNF4

Continue until Mid>Hi. 

Below is the implementation of above algorithm.

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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;
}
  
// This code is contributed by Shivi_Aggarwal

chevron_right


C

filter_none

edit
close

play_arrow

link
brightness_4
code

// C program to sort an array with 0, 1 and 2
// in a single pass
#include <stdio.h>
  
/* Function to swap *a and *b */
void swap(int* a, int* b);
  
// 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;
  
    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[0]);
    int i;
  
    sort012(arr, arr_size);
  
    printf("array after segregation ");
    printArray(arr, arr_size);
  
    getchar();
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java program to sort an array of 0, 1 and 2
import java.io.*;
  
class countzot {
  
    // Sort the input array, the array is assumed to
    // have values in {0, 1, 2}
    static void sort012(int a[], int arr_size)
    {
        int lo = 0;
        int hi = arr_size - 1;
        int mid = 0, temp = 0;
        while (mid <= hi) {
            switch (a[mid]) {
            case 0: {
                temp = a[lo];
                a[lo] = a[mid];
                a[mid] = temp;
                lo++;
                mid++;
                break;
            }
            case 1:
                mid++;
                break;
            case 2: {
                temp = a[mid];
                a[mid] = a[hi];
                a[hi] = temp;
                hi--;
                break;
            }
            }
        }
    }
  
    /* Utility function to print array arr[] */
    static void printArray(int arr[], int arr_size)
    {
        int i;
        for (i = 0; i < arr_size; i++)
            System.out.print(arr[i] + " ");
        System.out.println("");
    }
  
    /*Driver function to check for above functions*/
    public static void main(String[] args)
    {
        int arr[] = { 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 };
        int arr_size = arr.length;
        sort012(arr, arr_size);
        System.out.println("Array after seggregation ");
        printArray(arr, arr_size);
    }
}
/*This code is contributed by Devesh Agrawal*/

chevron_right


Python

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python program to sort an array with 0, 1 and 2 in a siingle pass
  
# Function to sort array
def sort012( a, arr_size):
    lo = 0
    hi = arr_size - 1
    mid = 0
    while mid <= hi:
        if a[mid] == 0:
            a[lo], a[mid] = a[mid], a[lo]
            lo = lo + 1
            mid = mid + 1
        elif a[mid] == 1:
            mid = mid + 1
        else:
            a[mid], a[hi] = a[hi], a[mid] 
            hi = hi - 1
return a
      
# Function to print array
def printArray( a):
    for k in a:
        print k,
    print
      
  
# Driver Program
arr = [0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1]
arr_size = len(arr)
arr = sort012( arr, arr_size)
print "Array aftter segregation :",
printArray(arr)
  
# Contributed by Harshit Agrawal

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# program to sort an
// array of 0, 1 and 2
using System;
  
class GFG {
    // Sort the input array, the array is assumed to
    // have values in {0, 1, 2}
    static void sort012(int[] a, int arr_size)
    {
        int lo = 0;
        int hi = arr_size - 1;
        int mid = 0, temp = 0;
  
        while (mid <= hi) {
            switch (a[mid]) {
            case 0: {
                temp = a[lo];
                a[lo] = a[mid];
                a[mid] = temp;
                lo++;
                mid++;
                break;
            }
            case 1:
                mid++;
                break;
            case 2: {
                temp = a[mid];
                a[mid] = a[hi];
                a[hi] = temp;
                hi--;
                break;
            }
            }
        }
    }
  
    /* Utility function to print array arr[] */
    static void printArray(int[] arr, int arr_size)
    {
        int i;
  
        for (i = 0; i < arr_size; i++)
            Console.Write(arr[i] + " ");
        Console.WriteLine("");
    }
  
    /*Driver function to check for above functions*/
    public static void Main()
    {
        int[] arr = { 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 };
        int arr_size = arr.Length;
        sort012(arr, arr_size);
  
        Console.Write("Array after seggregation ");
  
        printArray(arr, arr_size);
    }
}
  
// This code is contributed by Sam007

chevron_right


PHP

filter_none

edit
close

play_arrow

link
brightness_4
code

<?php 
// PHP program to sort an array
// with 0, 1 and 2 in a single pass
  
// Sort the input array, the array is 
// assumed to have values in {0, 1, 2}
function sort012(&$a, $arr_size)
{
    $lo = 0;
    $hi = $arr_size - 1;
    $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 */
function swap(&$a, &$b)
{
    $temp = $a;
    $a = $b;
    $b = $temp;
}
  
/* Utility function to print array arr[] */
function printArray(&$arr, $arr_size)
{
    for ($i = 0; $i < $arr_size; $i++)
        echo $arr[$i]." ";
    echo "\n";
}
  
// Driver Code
$arr = array(0, 1, 1, 0, 1, 2, 
             1, 2, 0, 0, 0, 1);
$arr_size = sizeof($arr);
  
sort012($arr, $arr_size);
  
echo "array after segregation ";
printArray($arr, $arr_size);
  
// This code is contributed
// by ChitraNayal
?>

chevron_right


Output:

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

Time Complexity: O(n)


The above code performs unnecessary swaps for some inputs which are not really required. We can modify the swap function to do a check that the values being swapped are same or not, if not same, then only swap the values. Thanks to Ankur Roy for suggesting this optimization.

Another approach: Count the number of 0s, 1s and 2s in the given array. Then store all the 0s in the beginning followed by all the 1s then all the 2s.

Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;
  
// Utility function to print the contents of an array
void printArr(int arr[], int n)
{
    for (int i = 0; i < n; i++)
        cout << arr[i] << " ";
}
  
// Function to sort the array of 0s, 1s and 2s
void sortArr(int arr[], int n)
{
    int i, cnt0 = 0, cnt1 = 0, cnt2 = 0;
  
    // Count the number of 0s, 1s and 2s in the array
    for (i = 0; i < n; i++) {
        switch (arr[i]) {
        case 0:
            cnt0++;
            break;
        case 1:
            cnt1++;
            break;
        case 2:
            cnt2++;
            break;
        }
    }
  
    // Update the array
    i = 0;
  
    // Store all the 0s in the beginning
    while (cnt0 > 0) {
        arr[i++] = 0;
        cnt0--;
    }
  
    // Then all the 1s
    while (cnt1 > 0) {
        arr[i++] = 1;
        cnt1--;
    }
  
    // Finally all the 2s
    while (cnt2 > 0) {
        arr[i++] = 2;
        cnt2--;
    }
  
    // Print the sorted array
    printArr(arr, n);
}
  
// Driver code
int main()
{
    int arr[] = { 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 };
    int n = sizeof(arr) / sizeof(int);
  
    sortArr(arr, n);
  
    return 0;
}

chevron_right


Output:

0 0 0 0 0 1 1 1 1 1 2 2

Time Complexity: O(n)



My Personal Notes arrow_drop_up