How to Sort an array of 0s, 1s and 2s?

Problem: 

Given an array of size N containing only 0s, 1s, and 2s; sort the array in ascending order.

Example 1:

Input: 

N = 5
arr[]= {0 2 1 2 0}

Output:

0 0 1 2 2

Explanation:

0s 1s and 2s are segregated 
into ascending order.

Example 2:

Input: 

N = 3
arr[] = {0 1 0}

Output:

0 0 1

Explanation:

0s 1s and 2s are segregated 
into ascending order.

Your Task:

You don’t need to read input or print anything. Your task is to complete the function sort012() that takes an array arr and N as input parameters and sorts the array in-place.

Expected Time Complexity: O(N)

Expected Auxiliary Space: O(1)

Solution

Java 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 segregation "); 
		printArray(arr, arr_size); 
	} 
} 

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

Leave a Comment

Your email address will not be published. Required fields are marked *