Sort an array of 0s, 1s and 2s | Dutch National Flag problem

Given an array A[] consisting of only 0s, 1s, and 2s. The task is to sort the array, i.e., put all 0s first, then all 1s and all 2s in last.

This problem is also the same as the famous “Dutch National Flag problem”. The problem was proposed by Edsger Dijkstra. The problem is as follows:

Given N balls of colour red, white or blue arranged in a line in random order. You have to arrange all the balls such that the balls with the same colours are adjacent with the order of the balls, with the order of the colours being red, white and blue (i.e., all red coloured balls come first then the white coloured balls and then the blue coloured balls). 

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}

Recommended Practice
Sort an array of 0s, 1s and 2s
Try It!

A Naive Solution solution is to simply sort the array using a standard sorting algorithm or sort library function. This solution requires O(n Log n) time.

A Better Solution is to traverse the array once and count number of 0s, 1s and 2s. Let these counts be c0, c1 and c2. Now traverse the array again, put c0 (count of 0s) 0s first, then c1 1s and finally c2 2s. This solution works in O(n) time, but this solution is not stable and requires two traversals of the array.

Sort an array of 0s, 1s, and 2s in one traversal and O(1) auxiliary space

This approach is based on the following idea:

  • The problem is similar to “Segregate 0s and 1s in an array”.
  • The problem was posed with three colors, here `0′, `1′ and `2′. The array is divided into four sections: 
    • arr[1] to arr[low – 1]
    • arr[low] to arr[mid – 1]
    • arr[mid] to arr[high – 1]
    • arr[high] to arr[n]
  • If the ith element is 0 then swap the element to the low range.
  • Similarly, if the element is 1 then keep it as it is.
  • If the element is 2 then swap it with an element in high range.

Illustration:

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

lo = 0, mid = 0, hi = 5

Step – 1: arr[mid] == 0

  • swap(arr[lo], arr[mid])
  • lo = lo + 1 = 1
  • mid = mid + 1 = 1
  • arr[] = {0, 1, 2, 0, 1, 2}

Step – 2: arr[mid] == 1

  • mid = mid + 1 = 2
  • arr[] = {0, 1, 2, 0, 1, 2}

Step – 3: arr[mid] == 2

  • swap(arr[mid], arr[hi])
  • hi = hi – 1 = 4
  • arr[] = {0, 1, 2, 0, 1, 2}

Step – 4: arr[mid] == 2

  • swap(arr[mid], arr[hi])
  • hi = hi – 1 = 3
  • arr[] = {0, 1, 1, 0, 2, 2}

Step – 5: arr[mid] == 1

  • mid = mid + 1 = 3
  • arr[] = {0, 1, 1, 0, 2, 2}

Step – 6: arr[mid] == 0

  • swap(arr[lo], arr[mid])
  • lo = lo + 1 = 2
  • mid = mid + 1 = 4
  • arr[] = {0, 0, 1, 1, 2, 2}

Hence, arr[] = {0, 0, 1, 1, 2, 2}

Follow the steps below to solve the given problem:

  • 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).
  • Traverse the array from start to end and mid is less than high. (Loop counter is i)
  • If the element is 0 then swap the element with the element at index low and update low = low + 1 and mid = mid + 1
  • If the element is 1 then update mid = mid + 1
  • 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
  • Print the array.
C++
// 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);

    printArray(arr, n);

    return 0;
}

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

/* 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]);
}

/* 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);

    printArray(arr, arr_size);

    getchar();
    return 0;
}
Java
// 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;
        // Iterate till all the elements
        // are sorted
        while (mid <= hi) {
            switch (a[mid]) {
                // If the element is 0
            case 0: {
                temp = a[lo];
                a[lo] = a[mid];
                a[mid] = temp;
                lo++;
                mid++;
                break;
            }
                // If the element is 1
            case 1:
                mid++;
                break;
                // If the element is 2
            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);
        printArray(arr, arr_size);
    }
}
/*This code is contributed by Devesh Agrawal*/
Python
# Python program to sort an array with
# 0, 1 and 2 in a single pass

# Function to sort array


def sort012(a, arr_size):
    lo = 0
    hi = arr_size - 1
    mid = 0
    # Iterate till all the elements
    # are sorted
    while mid <= hi:
        # If the element is 0
        if a[mid] == 0:
            a[lo], a[mid] = a[mid], a[lo]
            lo = lo + 1
            mid = mid + 1
        # If the element is 1
        elif a[mid] == 1:
            mid = mid + 1
        # If the element is 2
        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, end=' ')


# Driver Program
arr = [0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1]
arr_size = len(arr)
arr = sort012(arr, arr_size)
printArray(arr)

# Contributed by Harshit Agrawal
C#
// 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;
        // Iterate till all the elements
        // are sorted
        while (mid <= hi) {
            switch (a[mid]) {
            // If the element is 0
            case 0: {
                temp = a[lo];
                a[lo] = a[mid];
                a[mid] = temp;
                lo++;
                mid++;
                break;
            }
            // If the element is 1
            case 1:
                mid++;
                break;
            // If the element is 2
            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);

        printArray(arr, arr_size);
    }
}

// This code is contributed by Sam007
Javascript
<script>
// Javascript program to sort an array of 0, 1 and 2 

    // Sort the input array, the array is assumed to 
    // have values in {0, 1, 2} 
    function sort012(a,arr_size)
    {
        
        let lo = 0; 
        let hi = arr_size - 1; 
        let mid = 0;
        let temp = 0; 
        // Iterate till all the elements
        // are sorted
        while (mid <= hi)
        {
            // If the element is 0
            if(a[mid] == 0)
            {
                temp = a[lo]; 
                a[lo] = a[mid]; 
                a[mid] = temp; 
                lo++; 
                mid++; 
            }
            // If the element is 1
            else if(a[mid] == 1)
            {
                mid++; 
            }
            // If the element is 2
            else
            {
                temp = a[mid]; 
                a[mid] = a[hi]; 
                a[hi] = temp; 
                hi--;
            }
            
        }
    }
    
    /* Utility function to print array arr[] */
    function printArray(arr,arr_size)
    {
        let i;
        for (i = 0; i < arr_size; i++)
        {
            document.write(arr[i] + " ");
        }
        document.write("<br>");
    }
    
    /*Driver function to check for above functions*/
    let arr= [0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1 ];
    
    let arr_size = arr.length;
    sort012(arr, arr_size);
    printArray(arr, arr_size); 
    
    // This code is contributed by rag2127
</script>
PHP
<?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;
    // 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;
        }
    }
}

/* 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);

printArray($arr, $arr_size);

// This code is contributed
// by ChitraNayal
?>

Output
0 0 0 0 0 1 1 1 1 1 2 2 

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



Contact Us