Sort an Array of 0s 1s and 2s

Problem Statement

We are given an array consisting of only 0s, 1s, and 2s. Our task is to write a program to sort the given array without using inbuilt functions.

Example

Input: [0, 2, 1, 1, 2, 0]
Output: [0, 0, 1, 1, 2, 2]

Input: [0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0]
Output: [0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2] 

Input: [0]
Output: [0]

Example Explanation

Since we already know that an array is a linear data structure that contains data elements of the same data type at contiguous memory locations. For example, an array containing 12 elements as follows –[0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0].

Sorting an array means arranging of data elements in ascending or descending order, so our task here is to sort it in ascending order. Therefore, in our resultant sorted array, all 0s must come before all the 1s and all 1s must come before all 2s.

Thus, the sorted array will be [0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2]

Constraints

1 <= n <= 10^6 , here n is the number of elements in the array A 0 <= A[i] <= 2 , here A[i] is the value of every ith element in the array A

Approach 1: By Brute force method using double traversal

• Intuition: The first idea that comes into our mind since there are only 3 distinct values, is to just count the number of 0s, 1s and 2s. So we will be using 3 different variables to store these values which will be followed by overwriting the array with the help of these count variables.This problem is similar to three-way partitioning for the Dutch national flag problem.
• Algorithm:
  1. Initialize 3 variables count0, count1, count2 to store the count of 0s, 1s, and 2s respectively.
  2. Traverse through the array and increase count0 if the element is 0, increase count1 if the element is 1 and increase count2 if the element is 2.
  3. Now traverse the array again and overwrite the first count0 positions in the array with 0, the next count1 positions with 1, and the remaining count2 with 2.
• Implementation
  C++ Program to sort an array of 0s 1s and 2s using Brute Force method:

#include <bits/stdc++.h>
using namespace std;

// Function to sort 0s 1s and 2s
void sort_arr(int A[], int n)
{
    // Declaring 3 variables to store count
    int i, count0 = 0, count1 = 0, count2 = 0;

    // Traverse the complete array
    for (i = 0; i < n; i++) 
    {
        // If the element is 0
        if (A[i] == 0)
            count0 = count0 + 1;

        // If the element is 1
        else if (A[i] == 1)
            count1 = count1 + 1;

        // If the element is 2
        else
            count2 = count2 + 1;
    }

    // Updating the array
    i = 0;

    // Store all the 0s
    while (count0 > 0) {
        A[i++] = 0;
        count0--;
    }

    // Store all the 1s
    while (count1 > 0) {
        A[i++] = 1;
        count1--;
    }

    // Store all the 2s
    while (count2 > 0) {
        A[i++] = 2;
        count2--;
    }
}

int main()
{
    int arr[] = {0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0};
    int n = sizeof(arr) / sizeof(int);

    // Calling the function
    sort_arr(arr, n);

    // Printing the sorted array
    for (int i = 0; i < n; i++)
        cout << arr[i] << " ";

    return 0;
}

Java Program to sort an array of 0s 1s and 2s using Brute Force method:

import java.io.*;
class Main {

    // Function to sort the array of 0s 1s and 2s
    static void sort_arr(int A[], int n)
    {
        // Declaring 3 variables to store count
        int i = 0, count0 = 0, count1 = 0, count2 = 0;

        // Traverse the array completely
        while ( i < n )  
        {  
            // If the element is 0
            if ( A[i] == 0 )  
            {  
                count0 = count0 + 1;  
            }  
            // If the element is 1
            if ( A[i] == 1 )  
            {  
                count1 = count1 + 1;  
            }  
            // If the element is 2
            if ( A[i] == 2 )  
            {  
                count2 = count2 + 1;  
            }  

            i = i + 1;  
        }  

        // Updating the array
        i = 0;

        // Store all the 0s 
        while (count0 > 0) {
            A[i++] = 0;
            count0--;
        }
        // Store all the 1s
        while (count1 > 0) {
            A[i++] = 1;
            count1--;
        }
        // Store all the 2s
        while (count2 > 0) {
            A[i++] = 2;
            count2--;
        }
    }

    public static void main(String[] args)
    {
        int arr[] = {0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0};
        int n = arr.length;
        
        // Calling the function
        sort_arr(arr, n);
        
        // Printing the array
        for (int i = 0; i < n; i++)
            System.out.print(arr[i] + " ");
    }
}

Python Program to sort an array of 0s 1s and 2s using the Brute Force method:

def sort_arr(A, n):
    # Declaring 3 variables to store count
    count0 = 0
    count1 = 0
    count2 = 0

    # Traverse the complete array
    for i in range(n):
        if A[i] == 0:
            count0+=1

        elif A[i] == 1:
            count1+=1

        elif A[i] == 2:
            count2+=1

    # Updating the array
    i = 0

    # Store all the 0s 1s and 2s in the array
    while (count0 > 0):
        A[i] = 0
        i+=1
        count0-=1

    while (count1 > 0):
        A[i] = 1
        i+=1
        count1-=1

    while (count2 > 0):
        A[i] = 2
        i+=1
        count2-=1

arr = [0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0]
n = len(arr)

# Calling the sort_arr function
sort_arr(arr, n)

# Printing the sorted array
print(arr)

Output

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

Complexity Analysis
In the above approach, we require traversing the array twice; once to count the number of 0s, 1s, and 2s. While the second traversal is required to store the values at their sorted places in the array, that means we need to copy all the n elements. T(n) = O(n) + O(n)Hence, Time complexity will be – T(n) = O(n)Also, we observe that no extra space is required to store the elements.Hence, Space complexity will be – S(n) = O(1)

Approach 2: By using a single traversal method named the three-way partitioning

• Intuition: In this approach we will be using 3-pointers low, mid, high. We will use these pointers to place all 0s to the left end, all 2s to the right end, and all 1s in the middle of the array, hence making the array sorted.
• Algorithm:
  1. Initialize 3 pointers low = 0, mid = 0 and high = N-1.
  2. Start traversing the array from start to end, repeat this process until mid <= high.
     • If the A[mid]==0, then swap(A[low], A[mid]) and update the pointers low = low + 1 and mid = mid + 1.
     • If the A[mid]==1, just update the pointer mid = mid + 1.
     • If the A[mid]==2, then swap(A[mid], A[high]) and update the pointer high = high – 1
  3. Print the sorted array

• Implementation:
  C++ Program to sort an array of 0s 1s and 2s using three-way partitioning:

#include <bits/stdc++.h>
using namespace std;

// Three Way Partitioning function
void sort_3way(int A[], int n)
{
    // Declaring 3 pointers         
    int low = 0;
    int mid = 0;
    int high = n - 1;

    // Traverse array till completely sorted
    while (mid <= high) {
        switch (A[mid]) {

        // If the element is 0
        case 0:
            // Swapping and updating pointers
            swap(A[low++], A[mid++]);
            break;

        // If the element is 1 .
        case 1:
            mid++;
            break;

        // If the element is 2
        case 2:
            swap(A[mid], A[high--]);
            break;
        }
    }
}

int main()
{
    int arr[] = {0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0};
    int n = sizeof(arr) / sizeof(int);

    // Calling the sort function
    sort_3way(arr, n);

    // Printing the sorted array
    for (int i = 0; i < n; i++)
        cout << arr[i] << " ";

    return 0;
}

Java Program to sort an array of 0s 1s and 2s using three-way partitioning:

import java.io.*;

class Main
{
    // Three Way Partitioning function
    public static void sort_3way(int[] A, int n)
    {
        // Declaring 3 pointers
        int low = 0;
        int mid = 0;
        int high = n - 1;

        // Traverse array till completely sorted
        while (mid <= high)
        {
            // If the element is 0
            if (A[mid] == 0)     
            {
                swap(A, low, mid);
                ++low;
                ++mid;
            }

            // If the element is 1
            else if (A[mid] == 1)
            {
                ++mid;
            }

            // If the element is 0
            else { 
                swap(A, mid, high);
                --high;
            }
        }
    }

    // Swapping function
    private static void swap(int[] A, int i, int j)
    {
        int temp = A[i];
        A[i] = A[j];
        A[j] = temp;
    }

    public static void main (String[] args)
    {
        int arr[] = {0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0};
        int n = arr.length;

        // Calling the function
        sort_3way(arr, n);

        // Printing the sorted array 
        for (int i = 0; i < n; i++)
            System.out.print(arr[i] + " ");
    }
}

Python Program to sort an array of 0s 1s and 2s using three-way partitioning:

def sort_3way(A, n):

    # Declaring 3 pointers
    low = 0
    mid = 0
    high = n - 1

    # Traverse the array till completely sorted
    while (mid <= high):
        
        # If element is 0
        if A[mid] == 0:
            # Swapping
            A[low], A[mid] = A[mid], A[low]
            # Updating pointers
            low = low + 1
            mid = mid + 1

        # If element is 1
        elif A[mid] == 1:
            mid = mid + 1
        
        # If element is 2
        else:
            A[mid], A[high] = A[high], A[mid]
            high = high - 1

    return A

arr = [0, 1, 2, 2, 1, 0, 0, 2, 0, 1, 1, 0]
n = len(arr)

# Calling the sort function
sort_3way( arr, n)

# printing the sorted array
print(arr)

Output

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

• Complexity Analysis
  In this approach, we require traversing the array only once for comparing the value present at mid index with 0, 1, and 2 and swapping them if necessary. After the traversal is done, the result will be a sorted array. T (n) = C + O (n)
  Here C is the constant time taken for swapping the values in the variables.Hence Time complexity will be – T(n) = O(n)
  Also, we observe that no extra space is required to store the elements.Hence Space complexity will be – S(n) = O(1)

Conclusion

• In the first approach, we used a simple solution using the counting sort algorithm by counting the number of 0s, 1s, and 2s and then placing them in the array in their correct order.
• In the second approach, since we had only three distinct elements to be sorted, we divided the given array into four sections using three-pointers namely low, mid, and high.
• This problem of sorting 0s 1s and 2s is similar to three-way partitioning for the Dutch national flag problem.
• Comparison of Time and Space complexities for both approaches
  • Brute force approach: Time = O(n), Space = O(1).
  • Three-way partitioning: Time = O(n), Space = O(1).