Sort an array of 0s, 1s and 2s (Simple Counting)

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.

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}

Count the number of 0’s, 1’s and 2’s. After Counting, put all 0’s first, then 1’s and lastly 2’s in the array. We traverse the array two times. Time complexity will be O(n). 

C++

// Simple C++ program to sort an array of 0s 1s and 2s.
#include <iostream>
using namespace std;

void sort012(int* arr, int n)
{
    // Variables to maintain the count of 0's, 1's and 2's
    // in the array
    int count0 = 0, count1 = 0, count2 = 0;
    for (int i = 0; i < n; i++) {
        if (arr[i] == 0)
            count0++;
        if (arr[i] == 1)
            count1++;
        if (arr[i] == 2)
            count2++;
    }

    // Putting the 0's in the array in starting.
    for (int i = 0; i < count0; i++)
        arr[i] = 0;

    // Putting the 1's in the array after the 0's.
    for (int i = count0; i < (count0 + count1); i++)
        arr[i] = 1;

    // Putting the 2's in the array after the 1's
    for (int i = (count0 + count1); i < n; i++)
        arr[i] = 2;

    return;
}

// Prints the array
void printArray(int* arr, int n)
{
    for (int i = 0; i < n; i++)
        cout << arr[i] << " ";
    cout << endl;
}

// 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);
    printArray(arr, n);
    return 0;
}

// This code is contributed by Aditya Kumar (adityakumar129)

C

// Simple C program to sort an array of 0s 1s and 2s.
#include <stdio.h>

void sort012(int* arr, int n)
{
    // Variables to maintain the count of 0's, 1's and 2's
    // in the array
    int count0 = 0, count1 = 0, count2 = 0;
    for (int i = 0; i < n; i++) {
        if (arr[i] == 0)
            count0++;
        if (arr[i] == 1)
            count1++;
        if (arr[i] == 2)
            count2++;
    }

    // Putting the 0's in the array in starting.
    for (int i = 0; i < count0; i++)
        arr[i] = 0;

    // Putting the 1's in the array after the 0's.
    for (int i = count0; i < (count0 + count1); i++)
        arr[i] = 1;

    // Putting the 2's in the array after the 1's
    for (int i = (count0 + count1); i < n; i++)
        arr[i] = 2;

    return;
}

// Prints the array
void printArray(int* arr, int n)
{
    for (int i = 0; i < n; i++)
        printf("%d ", arr[i]);
    printf("\n");
}

// 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);
    printArray(arr, n);
    return 0;
}

// This code is contributed by Aditya Kumar (adityakumar129)

Java

// Simple Java program to sort an array of 0s 1s and 2s.
import java.lang.*;
import java.util.*;

public class GfG {

    public static void sort012(int arr[], int n)
    {
        // Variables to maintain the count of 0's, 1's and
        // 2's in the array
        int count0 = 0, count1 = 0;
        int count2 = 0;
        for (int i = 0; i < n; i++) {
            if (arr[i] == 0)
                count0++;
            if (arr[i] == 1)
                count1++;
            if (arr[i] == 2)
                count2++;
        }

        // Putting the 0's in the array in starting.
        for (int i = 0; i < count0; i++)
            arr[i] = 0;

        // Putting the 1's in the array after the 0's.
        for (int i = count0; i < (count0 + count1); i++)
            arr[i] = 1;

        // Putting the 2's in the array after the 1's
        for (int i = (count0 + count1); i < n; i++)
            arr[i] = 2;

        printArray(arr, n);
    }

    // Prints the array
    public static void printArray(int arr[], int n)
    {
        for (int i = 0; i < n; i++)
            System.out.print(arr[i] + " ");
        System.out.println();
    }

    // Driver function
    public static void main(String args[])
    {

        int arr[] = { 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 };
        int n = 12;
        sort012(arr, n);
    }
}

// This code is contributed by Aditya Kumar (adityakumar129)

Python3

# Python C++ program to sort an array of 0s
# 1s and 2s.
import math

def sort012(arr, n):

    # Variables to maintain the count of 0's,
    # 1's and 2's in the array
    count0 = 0
    count1 = 0
    count2 = 0
    for i in range(0,n):
        if (arr[i] == 0):
            count0=count0+1
        if (arr[i] == 1):
            count1=count1+1
        if (arr[i] == 2):
            count2=count2+1


    # Putting the 0's in the array in starting.
    for i in range(0,count0):
        arr[i] = 0

    # Putting the 1's in the array after the 0's.
    for i in range( count0, (count0 + count1)) :
        arr[i] = 1

    # Putting the 2's in the array after the 1's
    for i in range((count0 + count1),n) :
        arr[i] = 2

    return


# Prints the array
def printArray( arr,  n):

    for i in range(0,n):
        print( arr[i] , end=" ")
    print()


# Driver code
arr = [ 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 ]
n = len(arr)
sort012(arr, n)
printArray(arr, n)

# This code is contributed by Gitanjali.

C#

// Simple C# program
// to sort an array of 0s
// 1s and 2s.
using System;

public class GfG{

    public static void sort012(int []arr, int n)
    {
        // Variables to maintain
        // the count of 0's,
        // 1's and 2's in the array
        int count0 = 0, count1 = 0;
        int count2 = 0;
        for (int i = 0; i < n; i++) {
            if (arr[i] == 0)
                count0++;
            if (arr[i] == 1)
                count1++;
            if (arr[i] == 2)
                count2++;
        }

        // Putting the 0's in the
        // array in starting.
        for (int i = 0; i < count0; i++)
            arr[i] = 0;

        // Putting the 1's in the
        // array after the 0's.
        for (int i = count0; i <
            (count0 + count1); i++)
            arr[i] = 1;

        // Putting the 2's in the
        // array after the 1's
        for (int i = (count0 + count1);
            i < n; i++)
            arr[i] = 2;

        printArray(arr, n);
    }

    // Prints the array
    public static void printArray(int []arr, int n)
    {
        for (int i = 0; i < n; i++)
            Console.Write(arr[i] + " ");
        Console.WriteLine();
    }

    // Driver function
    public static void Main()
    {

        int []arr = { 0, 1, 1, 0, 1, 2, 1,
                    2, 0, 0, 0, 1 };
        int n = 12;
        sort012(arr, n);
    }
}

// This code is contributed by vt_m

Javascript

<script>

// JavaScript program
// to sort an array of 0s
// 1s and 2s.

    function sort012(arr, n)
    {
        // Variables to maintain
        // the count of 0's,
        // 1's and 2's in the array
        let count0 = 0, count1 = 0;
        let count2 = 0;
        for (let i = 0; i < n; i++) {
            if (arr[i] == 0)
                count0++;
            if (arr[i] == 1)
                count1++;
            if (arr[i] == 2)
                count2++;
        }

        // Putting the 0's in the
        // array in starting.
        for (let i = 0; i < count0; i++)
            arr[i] = 0;

        // Putting the 1's in the
        // array after the 0's.
        for (let i = count0; i <
            (count0 + count1); i++)
            arr[i] = 1;

        // Putting the 2's in the
        // array after the 1's
        for (let i = (count0 + count1);
            i < n; i++)
            arr[i] = 2;

        printArray(arr, n);
    }

    // Prints the array
    function printArray(arr, n)
    {
        for (let i = 0; i < n; i++)
            document.write(arr[i] + " ");
        document.write();
    }



// Driver code

        let arr = [ 0, 1, 1, 0, 1, 2, 1,
                    2, 0, 0, 0, 1 ];
        let n = 12;
        sort012(arr, n);

</script>

Output

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

Time Complexity: O(n)

Auxiliary Space: O(1)
Problems with the above solution.:

1. It requires two traversals of array.
2. This solution may not work if values are a part of the structure. For example, consider a situation where 0 represents Computer Science Stream, 1 represents Electronics and 2 represents Mechanical. We have a list of student objects (or structures) and we want to sort them. We cannot use the above sort as we simply put 0s, 1s, and 2s one by one.

Another Approach:

C++

// C++ program
#include <bits/stdc++.h>
using namespace std;
// Example
//
// input = [0, 1, 2, 2, 0, 0]
// output = [0, 0, 0, 1, 2, 2]

int main() {

    vector<int> inputArray={0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1};
    vector<int> outputArray;
    int indexOfOne = 0;
    for(int item: inputArray)
    {
        if(item==2)
        {   outputArray.push_back(item);
        }
        else if(item==1)
        {   outputArray.insert(outputArray.begin() + indexOfOne, item);
            indexOfOne+=1;
        }
        else if(item==0)
        {   outputArray.insert(outputArray.begin(), item);
            indexOfOne+=1;
        }
        else
        {   cout<<" wrong value - Aborting ";
            continue;
        }
    }

    for(int i=0;i<outputArray.size();i++)
    {
        cout<<outputArray[i]<<" ";
    }

    return 0;
}
// This code is contributed by Pushpesh Raj

Java

import java.util.ArrayList;
import java.util.List;

// Example
//
// input = [0, 1, 2, 2, 0, 0]
// output = [0, 0, 0, 1, 2, 2]
class GFG {
    static int[] inputArray = { 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 };
    static List<Integer> outputArray = new ArrayList<>();
    static int indexOfOne = 0;

    static void print() {
        for (int item : inputArray)
            if (item == 2)
                outputArray.add(item);
            else if (item == 1) {
                outputArray.add(indexOfOne, item);
                indexOfOne += 1;
            } else if (item == 0) {
                outputArray.add(0, item);
                indexOfOne += 1;
            } else {
                System.out.println(" wrong value - Aborting ");
                continue;
            }
    }

    public static void main(String[] args) {
        print();
        for (int item : outputArray)
            System.out.print(item+", ");
    }
}
// This code is contributed by Amit Katiyar

Python3

# Example
#
# input = [0, 1, 2, 2, 0, 0]
# output = [0, 0, 0, 1, 2, 2]

inputArray = [0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1]
outputArray = []
indexOfOne = 0
for item in inputArray:
    if item == 2:
        outputArray.append(item)
    elif item == 1:
        outputArray.insert(indexOfOne, item)
        indexOfOne += 1
    elif item == 0:
        outputArray.insert(0, item)
        indexOfOne += 1
    else:
        print(" wrong value - Aborting ")
        continue
print(outputArray)

C#

using System;
using System.Collections.Generic;

// Example
//
// input = [0, 1, 2, 2, 0, 0]
// output = [0, 0, 0, 1, 2, 2]

class GFG{

static int[] inputArray = { 0, 1, 1, 0, 1, 2,
                            1, 2, 0, 0, 0, 1 };
static List<int> outputArray = new List<int>();
static int indexOfOne = 0;

static void print()
{
    foreach (int item in inputArray)
        if (item == 2)
            outputArray.Add(item);
        else if (item == 1)
        {
            outputArray.Insert(indexOfOne, item);
            indexOfOne += 1;
        }
        else if (item == 0)
        {
            outputArray.Insert(0, item);
            indexOfOne += 1;
        }
        else
        {
            Console.WriteLine(" wrong value - Aborting ");
            continue;
        }
}

// Driver code
public static void Main(String[] args)
{
    print();
    foreach(int item in outputArray)
        Console.Write(item + ", ");
}
}

// This code is contributed by 29AjayKumar

Javascript

<script>

// Example
//
// input = [0, 1, 2, 2, 0, 0]
// output = [0, 0, 0, 1, 2, 2]
let inputArray=[ 0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 ];
let outputArray = [];
let indexOfOne = 0;

function print()
{
    for (let item of inputArray.values())
            if (item == 2)
                outputArray.push(item);
            else if (item == 1) {
                outputArray.splice(indexOfOne,0, item);
                indexOfOne += 1;
            } else if (item == 0) {
                outputArray.splice(0,0, item);
                indexOfOne += 1;
            } else {
                document.write(" wrong value - Aborting ");
                continue;
            }
}

print();
for (let item of outputArray.values())
    document.write(item+", ");

// This code is contributed by rag2127
</script>

Output:

0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 

Time Complexity: O(n)

Auxiliary Space: O(1)

Optimal Solution that handles above issues : 
Sort an array of 0s, 1s and 2s (Dutch National Flag Algorithm)