Home » Data Structure » Array Data Structure

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

C++ implementation to sort an array of 0s, 1s and 2s.
Submitted by Vikneshwar GK, on February 05, 2022

Consider an array of size n and containing only 0s, 1s, and 2s. The task at hand is to sort this array, i.e., place all the 0s first followed by 1s and then 2s.

Example:

Input:
test1[] = {1, 1, 2, 2, 1, 0, 0, 0, 1, 2}

Output:
test1[] = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2} 

Input:
test1[] = {1, 0, 2, 2, 0, 1, 0, 0, 1, 0}

Output:
test1[] = {0, 0, 0, 0, 0, 1, 1, 1, 2, 2} 

Solution 1: Using Pointers

This method involves using three-pointers that handle the placement of 0s, 1s, and 2s. If the pointers are low, mid, and high, and the size of the array is N, then up to low indicate the 0s, low to mid indicate 1s, mid to high contains unknown elements, and high to N contains the 2s. The switching and swapping are based on the following algorithm:

• Initialize low=0, mid=0, and high=N-1
• Iterate through the array from 0 to N-1 using mid as the index while keeping the condition mid<=high
• If the index element is 0, then swap it with index low. Then increment low and mid
• If the index element is 1, then only increment mid
• If the index element is 2, then swap the element with high. Then decrement high
• Finally, print the array

C++ Implementation:

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

// Function to sort 0s,1s,2s
void sort(int array[], int length)
{
    int low = 0;
    int high = length - 1;
    int mid = 0;

    // iterate through the array
    // mid<=high
    while (mid <= high) {

        // mid is the iterating index
        switch (array[mid]) {

        case 0:
            swap(array[low++], array[mid++]);
            break;
        case 1:
            mid++;
            break;
        case 2:
            swap(array[mid], array[high--]);
            break;
        }
    }
}

// Function to print the array
void printArray(int array[], int length)
{
    for (int i = 0; i < length; i++)
        cout << array[i] << " ";
    cout << endl;
}

// Main function
int main()
{
    int array[100], N;
    
    cout << "Enter Number of elements: ";
    cin >> N;

    for (int i = 0; i < N; i++) {
        cout << "Enter element " << i + 1 << ":";
        cin >> array[i];
    }

    // sorting the array
    sort(array, array + N);

    cout << "Sorted Array" << endl;
    printArray(array, N);

    return 0;
}

Output:

Enter Number of elements: 10
Enter element 1:1
Enter element 2:1
Enter element 3:2
Enter element 4:0
Enter element 5:0
Enter element 6:1
Enter element 7:2
Enter element 8:1
Enter element 9:0
Enter element 10:2
Sorted Array
0 0 0 1 1 1 1 2 2 2

Time Complexity: O(n)

Solution 2: Using Counters

This is a simple method that involves three counters to count 0s, 1s, and 2s, respectively. The idea here is to count the number of 0s, 1s, and 2s, and place replace the array such many times of the respective element. It involves the following steps:

• Initialize c0, c1, and c2 to 0
• Traverse through the array and increment c0 if the element is 0, c1 if the element is 1, and c2 if the element is 2
• Now iterate through the array place first c0 elements as 0, next c1 elements as 1, and last c2 elements as 2

C++ Implementation:

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

// Function to sort 0s,1s,2s
void sort(int array[], int length)
{
    int c0 = 0, c1 = 0, c2 = 0;

    // Count the number of 0s, 1s and 2s in the array
    for (int i = 0; i < length; i++) {
        switch (array[i]) {
        case 0:
            c0++;
            break;
        case 1:
            c1++;
            break;
        case 2:
            c2++;
            break;
        }
    }

    int i = 0;
    // update 0s in the array
    while (c0 > 0) {
        array[i++] = 0;
        c0--;
    }

    // update 1s in the array
    while (c1 > 0) {
        array[i++] = 1;
        c1--;
    }

    // update 2s in the array
    while (c2 > 0) {
        array[i++] = 2;
        c2--;
    }
}

// Function to print the array
void printArray(int array[], int length)
{
    for (int i = 0; i < length; i++)
        cout << array[i] << " ";
    cout << endl;
}

// Main function
int main()
{
    int array[100], N;
    
    cout << "Enter Number of elements: ";
    cin >> N;

    for (int i = 0; i < N; i++) {
        cout << "Enter element " << i + 1 << ":";
        cin >> array[i];
    }

    // sorting the array
    sort(array, array + N);

    cout << "Sorted Array" << endl;
    printArray(array, N);

    return 0;
}

Output:

Enter Number of elements: 10
Enter element 1:2
Enter element 2:2
Enter element 3:1
Enter element 4:0
Enter element 5:0
Enter element 6:1
Enter element 7:2
Enter element 8:0
Enter element 9:0
Enter element 10:1
Sorted Array
0 0 0 0 1 1 1 2 2 2

Time Complexity: O(n)