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

Difficulty: Medium, Asked-in: Google, Facebook, Microsoft, Amazon, Paytm, Morgan Stanley, Hike, Walmart, Adobe, SAP Labs, Qualcomm.

Key takeaway

A special sorting problem that can be solved using O(n) time complexity.
There is a lot of variation available to this coding problem.
We are using a single loop and three-pointers to solve the problem efficiently.
This problem is a variation of the famous Dutch national flag problem.

Let’s understand the problem

Given an array X[] consisting of 0s, 1s, and 2s, write a program to sort the array of 0’s, 1’s, and 2’s in ascending order.

Examples

Input: X[] = [0, 2, 1, 0, 1, 2, 1, 0], Output: X[] = [0, 0, 0, 1, 1, 1, 2, 2].

Input: X[]= [0, 1, 1, 0, 1, 2, 1, 2, 0, 0], Output: X[] = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2].

Input: X[] = [2, 1, 0], Output: X[] = [0, 1, 2].

Discussed solution approaches

Brute force approach using double traversal
Efficient approach using single scan: Three-way partitioning

Brute force approach using double traversal

Solution idea and steps

The basic idea that comes to mind is to simply count the number of 0’s, 1’s, and 2’s. Then, place all 0’s at the beginning of the array followed by 1’s and 2's.

We declare variables countZero, countOne, countTwo to store the frequency of 0, 1, and 2.
Now we traverse the array to count the frequency of 0’s, 1’s, and 2's.
Now, we again traverse the array and replace the first countZero places of the array as 0, next countOne places by 1, and last countTwo places by 2.

Solution pseudocode

sort_0_1_2(X, n)
    countZero = 0, countOne = 0, countTwo = 0

    for i from 0 to n - 1
        if X[i] equals 0
            countZero = countZero + 1
        if X[i] equals 1
            countOne = countOne + 1
        if X[i] equals 2
            countTwo = countTwo + 1

    i = 0
    while countZero is not 0
        X[i] = 0
        i = i + 1
        countZero = countZero - 1

    while countOne is not 0
        X[i] = 1
        i = i + 1
        countOne = countOne - 1

    while countTwo is not 0
        X[i] = 2
        i = i + 1
        countTwo = countTwo - 1

Solution analysis

Time complexity = Time complexity of counting 0’s, 1’s, and 2’s + Time complexity of placing 0’s, 1’s, and 2’s in the sorted order = O(n) + O(n) = O(n). As we are using constant extra space, space complexity = O(1).

Efficient approach using single scan: Three-way partitioning

Solution idea

The critical question is: can we solve the problem in a single traversal? We have only three unique elements in the input, so can we use an idea similar to the partition process of the quicksort to place the elements at their correct position in the sorted output? Let's think!

We can solve the problem using a single scan by maintaining the correct order of 0’s, 1’s, and 2’s using variables. Actually, we have three types of elements to be placed in sorted order, so we divide the given array into four sections using three-pointers. Let us name these pointers as low, mid and high.

X[0 … low - 1] contains only 0s.
X[low … mid - 1] contains only 1s.
X[mid … high] contains elements we have not examined yet.
X[high + 1 … n - 1] contains only 2s.

Solution steps

We initialise three pointers low = 0, mid = 0 and high = n - 1.

Now we scan the array using the mid pointer. In other words, we run a loop until mid <= high.

If X[mid] == 0: Swap X[mid] with X[low], then increment both low and mid. The 0 is now in the 0s region and the unknown region shrinks by one.
If X[mid] == 1: Increment mid. The 1 is already in the correct region, so the unknown region shrinks by one.
If X[mid] == 2: Swap X[mid] with X[high], then decrement high. The 2 is now in the 2s region. We are keeping mid unchanged because the value swapped into X[mid] came from the unknown region and still needs to be checked.

Solution pseudocode

sort_0_1_2(X, n)
    low = 0, mid = 0, high = n - 1

    while mid <= high
        if X[mid] equals 0
            swap X[low] and X[mid]
            low = low + 1
            mid = mid + 1

        else if X[mid] equals 1
            mid = mid + 1

        else if X[mid] equals 2
            swap X[mid] and X[high]
            high = high - 1

Time and space complexity analysis

In the worst case, we will be traversing and accessing each element only once. Time complexity = O(n). Space complexity = O(1), as we are using a constant number of variables.

Critical ideas to think!

Can we optimize the above code and reduce the swap operations?
Is this algorithm stable? If not, then how to make it stable?
Why are we not increasing the mid-index if (X[mid] ==2)?
Can we solve the problem using the partition process of the quick sort?
Can we solve this problem using two-pointers?
How can we modify the above code for a similar problem to sort the array containing 0s,1s, 2s, and 3s?
What are the famous linear time sorting algorithms? Can we apply the idea of counting sort to solve the above problem?

Comparisons of time and space complexities

Brute force approach: Time = O(n), Space = O(1).
Three-way partitioning: Time = O(n), Space = O(1).

Sort 0s, 1s and 2s Without using Extra Space

Let’s understand the problem

Examples

Discussed solution approaches

Brute force approach using double traversal

Solution idea and steps

Solution pseudocode

Solution analysis

Efficient approach using single scan: Three-way partitioning

Solution idea

Solution steps

Solution pseudocode

Time and space complexity analysis

Critical ideas to think!

Comparisons of time and space complexities

Suggested coding questions to practice

More from EnjoyAlgorithms

Self-paced Courses and Blogs