[Problem Solving] Sort Colors – Variation of Dutch National Flag Problem

Given an array containing integers 0, 1 and 2 representing the colors red, white and blue, sort them in-place so that the same numbers are adjacent, with the order being 0, 1 and 2.

Example:

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

Solution #1:

There are only three different values in the array and those are 0, 1 and 2. We have to return an array arranged in an order with all the numbers with same values clubbed together. A simple approach is to count the number of 0s, 1s and 2s in the given array and overwrite the given array with these three values according to the counts. In the given array,

count(0) = 2

count(1) = 2

count(2) = 2

So let’s say we store these values in an array counts = [[0, 2], [1, 2], [2, 2]]. We can then iterate through this array and overwrite the given input array with 0, 1 and 2 according to their respective counts.

Time Complexity: O(N) to count the values and another O(N) to overwrite the input array, so the overall time complexity will be O(N).

Space Complexity: O(C) since we use constant space to store the values 0, 1, 2 and their counts.

Solution #2:

Notice that we have to shuffle around only three different values. This means that we can use a few pointers to keep track of some of the values.

Let’s use three different pointers as follows.

• zero_ptr to keep track of the leftmost index where we can overwrite with a 0.
• two_ptr to keep track of the rightmost index where we can overwrite with a 2.
• curr pointer to keep track of the current index we are looking at.

Algorithm to get the Dutch National Flag pattern:

• When curr points to a 0, we swap it with the element at zero_ptr. After swapping we move curr and zero_ptr one step forward.
• When curr points to a 2, we swap it with the element at two_ptr. We then decrement the value of two_ptr. Note that we should not move curr at this step because we may have just swapped a 2 with another 2 and we should re-process it in the next step.
• If curr points to a 1, we just move curr to the next element.

The above loop stops when curr > two_ptr. This indicates that we have placed all the values in their appropriate slots. If this is too confusing, look at the picture below to see how the algorithm works step by step.

Time Complexity: O(N)

Space Complexity: O(1) for in-place re-arrangement

Code:

def sortColors(self, nums: List[int]) -> None:
    """
    Do not return anything, modify nums in-place instead.
    """

    zero_ptr = curr = 0
    two_ptr = len(nums) - 1

    while curr <= two_ptr:
        if nums[curr] == 0:
            nums[curr], nums[zero_ptr] = nums[zero_ptr], nums[curr]
            zero_ptr += 1
            curr += 1
        elif nums[curr] == 2:
            nums[curr], nums[two_ptr] = nums[two_ptr], nums[curr]
            two_ptr -= 1
        else:  
            curr += 1

Hope you learnt how to solve this fun problem. Have a great day!