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.
• For the efficient solution, we are using a single loop and three-pointers to build the partial solution. This problem is a simple 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

• A brute force approach using double traversal
• An efficient approach using single scan: Three-way partitioning

A 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 three variables countZero, countOne, countTwo to store the frequency of 0’s, 1’s, and 2’s in the array.

  int countZero = 0, countOne = 0, countTwo = 0
  

• Now we traverse through the array to update the frequency of 0’s, 1’s, and 2's.

  for(int i = 0; i < n; i = i + 1)
  {
    if(X[i] == 0) countZero = countZero + 1
    if(X[i] == 1) countOne = countOne + 1
    if(X[i] == 2) countTwo = countTwo + 1   
  }
  

• 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.

  int i = 0
  while(countZero != 0)
  {
    X[i] = 0
    i = i + 1
    countZero = countZero - 1
  }
  
  while(countOne != 0)
  {
    X[i] = 1
    i = i + 1
    countOne = countOne - 1
  }
  
  while(countTwo != 0)
  {
    X[i] = 2
    i = i + 1
    countTwo = countTwo - 1
  }
  

Solution Pseudocode

Solution Analysis

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

An efficient approach using single scan: Three-way partitioning

Solution Idea

The critical question is: can we solve the problem in a single traversal of the array? 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, therefore, we divide the given array into four sections using three-pointers. Let us name these pointers as low, mid, and high.

1. X[0…low-1] to fill zeroes.
2. X[low..mid-1] to fill ones.
3. X[mid…high] unknown
4. X[high+1..n-1] only twos
5. Think of low and high as the range [low, high) for values in the middle partition. The middle partition will contain the values from low up to, but not including, high.
6. All values are less than low end up in the left partition.
7. Finally, all values greater than or equal to high end up in the right partition.

Solution Steps

• Initialise three pointers low = 0, mid = 0 and high = n -1

Run a loop until mid<=high

• If (X[mid] ==0), then swap the element with the element low and shrink the unknown range by doing low++ and mid++.
• If (X[mid] ==1), then simply increment mid (mid++), thus further shrinking unknown range.
• If (X[mid] ==2), then swap it with an element in the high range and decrement high (high--). Then, again traverse the array from mid-index only, as the element at this index has not been considered yet.

Solution Pseudocode

Time and space complexity analysis

In the worst case, we are traversing and accessing each element of the array 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)

Suggested coding questions to practice

• Move zeroes to the end of an array
• Partition algorithm in the quicksort
• Container with most water
• Trapping rainwater problem
• Check for pair in an array with a given sum
• Check whether an array is a subset of another array
• Remove duplicates from sorted array
• Equilibrium index of an array
• Find max and min in an array
• Valid mountain array
• Roman to Integer
• n Repeated element in 2n size array

Thanks to Navtosh Kumar for his contribution in creating the first draft. Please write comments if you find an error or you want to share more insights about the topic. Enjoy learning, Enjoy coding, Enjoy algorithms!