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

2.7

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.

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

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 C implementation of above algorithm.

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;

    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;
}

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

Python

# 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)

# Conribued by Harshit Agrawal

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 inputs like 0 0 0 0 1 1 1 2 2 2 2 2 : lo=4 and mid=7 and hi=11. In present code: first 7 exchanged with 11 and hi become 10 and mid is still pointing to 7. again same operation is till the mid <= hi. But it is really not required. We can change the swap function to do a check that the values being swapped are same or not, if not same, then only swap values. Thanks to Ankur Roy for suggesting this optimization. Source: http://www.csse.monash.edu.au/~lloyd/tildeAlgDS/Sort/Flag/

Please write comments if you find the above code/algorithm incorrect, or find better ways to solve the same problem.

GATE CS Corner    Company Wise Coding Practice

Please write to us at contribute@geeksforgeeks.org to report any issue with the above content.

Recommended Posts:



2.7 Average Difficulty : 2.7/5.0
Based on 191 vote(s)










Writing code in comment? Please use ide.geeksforgeeks.org, generate link and share the link here.