Dutch National Flag problem - Sort 0, 1, 2 in an array

As seen in the image above, the Dutch National Flag has three colors:

Red at the top
White in the middle
Blue at the bottom

Edsger Dijkstra proposed a problem where red, blue, and white balls are given in random order. The task is to arrange them or sort them in the order of the colors in the Dutch national flag. That is the balls of the same color are to be placed together. This task is to be executed in O(n) time complexity and without taking any extra space that is in constant space O(1).

Here, we will represent the balls as:

red ball as 0
white ball as 1
blue ball as 2

We will consider these balls as numbers given in an array in random order. Let's see how we can solve this problem:

We have an input array that contains the numbers(balls) 0s, 1s, and 2s only in random order. The length of the array can vary but it should contain just these three numbers that are 0, 1, and 2. Our task is to sort the input array in ascending order. Let's see the input and output before moving towards the approach:

Input: [2, 0, 1, 1, 0]

Output: [0, 0, 1, 1, 2]

Input: [1, 0, 2, 2, 0, 1, 1]

Output: [0, 0, 1, 1, 1, 2, 2]

So, there are three solutions to this problem:

Sorting
Simple Counting
Dutch National Flag Algorithm

Method 1: Sorting - O(nlogn)

The Brute-Force solution to this problem is simply using any of the sorting methods. We can use Merge Sort or Quick Sort and land up with the time complexity of O(nlogn).

Method 2: Simple counting - O(n)

We can optimize the above-discussed approach by a simple counting approach. Here, we just linearly traverse the input array and keep a count of the occurrences of 0s, 1s, and 2s. At the end of the traversal, we can get the individual count of 0s, 1s, and 2s. Suppose the input array is [1, 0, 2, 2, 0, 1, 1] Here there are two 0s, three 1s, and two 2s. Now simply we can run a loop two times and insert two 0s, then run a loop for three times and insert three 1s and lastly run a loop two times and insert two 2s. We end up with a sorted array which looks like: [0, 0, 1, 1, 1, 2, 2]. this approach requires time complexity of O(n).

Method 3: Dutch National Flag Algorithm

This is the optimal solution to this problem. Here we use three pointers to achieve three-way partitioning. Let me put it simply. We will use three pointers that are low, mid, and high. Initially, we point the low and mid pointers to the start that is the first element of the input array, and the high pointer to the end that is the last element of the input array.

So, basically this algorithm is based on the fact all the numbers from arr[0] to arr[low-1] will be 0s and all the numbers from arr[high+1] to arr[n-1] will be 2s and arr[low] to arr[mid-1] will be 1s

arr[0] to arr[low-1] ----------0
arr[high+1] to arr[n-1]------2
arr[low] to arr[mid-1]--------1

That is all the numbers to the left of low will be 0 and all the numbers to the right of high will be 1. How to achieve this? we will traverse the input array until mid<=high. We will check three conditions that is when arr[mid]==0, arr[mid]==1, and arr[mid]==2. See the piece of code below:

while mid <= high:
    if arr[mid] == 0:
       arr[low], arr[mid] = arr[mid], arr[low]
       low = low + 1
       mid = mid + 1
    elif arr[mid] == 1:
       mid = mid + 1
    else:
       arr[mid], arr[high] = arr[high], arr[mid]
       high = high - 1

The above piece of code can be described as:

if arr[mid] == 0:

swap the elements at low and mid

increment low

increment mid

elif arr[mid] == 1:

increment mid

else:

swap the elements at mid and high

decrement high

Let's see each step of it. In the below steps, l = low, m = mid and h = high.

So, as mid>high, the while loop stops, and as said earlier, all 0s are to the left of the low and all 1s are to the right of the high.

Now, let's see the implementation of the same in Python:

def dutchsort(arr, n):
    low = 0
    high = n - 1
    mid = 0
    while mid <= high:
        if arr[mid] == 0:
            arr[low], arr[mid] = arr[mid], arr[low]
            low = low + 1
            mid = mid + 1
        elif arr[mid] == 1:
            mid = mid + 1
        else:
            arr[mid], arr[high] = arr[high], arr[mid]
            high = high - 1
    return arr


# Driver Program
arr = [0, 1, 2, 0, 1, 2]
n = len(arr)
arr = dutchsort(arr, n)
print(arr)

[0, 0, 1, 1, 2, 2]

Dutch National Flag problem - Sort 0, 1, 2 in an array