Vivek's Blog

Today I solved the problem of sorting colors, or, to say, sorting an array of zeros, ones, and twos. This problem involved an array containing only zeros, ones, and twos, and we had to sort it so that the same numbers were adjacent to each other.

1. Brute force approach

In this approach, I thought of iterating through the array once to count the number of zeros, ones, and twos. Then we'll iterate through the array and place these numbers, along with their counts, adjacent to each other so that, in the end, we get a new array in which all the numbers are arranged in non-decreasing order.

class Solution {
public:
    // Function to sort the array containing only 0s, 1s and 2s
    void sortZeroOneTwo(vector<int>& nums) {
        // Initialize count variables for 0s, 1s, and 2s
        int count0 = 0, count1 = 0, count2 = 0;
 
        // Count the frequency of 0s, 1s, and 2s in the array
        for(int i = 0; i < nums.size(); i++) {
            if(nums[i] == 0) count0++;
            else if(nums[i] == 1) count1++;
            else count2++;
        }
 
        // Overwrite the array with the counted values
        int index = 0;
 
        // Fill with 0s
        while(count0--) {
            nums[index++] = 0;
        }
 
        // Fill with 1s
        while(count1--) {
            nums[index++] = 1;
        }
 
        // Fill with 2s
        while(count2--) {
            nums[index++] = 2;
        }
    }
};

2. The better approach

In this approach, I thought, let's not do that, and let's not just iterate through the array once again and place the elements just there. What we can do is make it simpler and put 3 counters to count the zeros, ones, and twos. When we get this number, we can change the same array and iterate through it and place it at its respective positions. What it will do is we do not need to create a new array, and the space complexity will be lower.

class Solution {
public:
    // Function to sort an array containing only 0s, 1s, and 2s
    void sortZeroOneTwo(vector<int>& nums) {
        // Count of 0s, 1s, and 2s
        int cnt0 = 0, cnt1 = 0, cnt2 = 0;
 
        // First pass: Count the number of 0s, 1s, and 2s
        for (int i = 0; i < nums.size(); i++) {
            if (nums[i] == 0) cnt0++;
            else if (nums[i] == 1) cnt1++;
            else cnt2++;
        }
 
        // Second pass: Fill the array with 0s, then 1s, then 2s
 
        // Fill the first 'cnt0' elements with 0
        for (int i = 0; i < cnt0; i++) {
            nums[i] = 0;
        }
 
        // Fill the next 'cnt1' elements with 1
        for (int i = cnt0; i < cnt0 + cnt1; i++) {
            nums[i] = 1;
        }
 
        // Fill the remaining elements with 2
        for (int i = cnt0 + cnt1; i < nums.size(); i++) {
            nums[i] = 2;
        }
    }
};

3. The most optimal approach

Was to use the Dutch national flag algorithm. In this algorithm, what we needed to do was to initialize 3 pointers:

• low, which is pointing to the low
• mid, which is initialized with 0, but it will point to the middle element of the array
• high, which is pointing to the last element of the array, n - 1 Ok, ay n is the size of the array. Here, what we needed to do is, if we encounter the middle element or the element at the middle index, if it is zero, then we can swap this element with the lowest element, like the element at the low index. Similarly, when we swap this, then we increment the index of low as well as mid. If we encounter one at the middle index or the mid index, then we do not do anything and just increment the count of middle because that is where it should be. If it is one, then it should be in the middle, and if we encounter the middle element or the element at the high index, then we will swap it with the element at the high index. When we swap the mid element or the element at the mid with the element at the high index, then we will decrement, or we will decrease the value of the high index, so that we can come to the values which are left in the array. We will not increment the mid index because after swapping the mid element, there must be zero or two, so we will not increment the count because this also needs processing. That's why we will not implement it. At the end, we will get the same array but sorted, and we won't need to create a new array. We will not have to iterate again and again, and the time complexity will be O(n). It will be the same for every approach that I shared, but for this, the space complexity will be much lower than the others.

class Solution {
public:
    // Function to sort array containing 0s, 1s, and 2s using Dutch National Flag Algorithm
    void sortZeroOneTwo(vector<int>& nums) {
        // Initialize three pointers: low, mid starting from 0, high from end of array
        int low = 0, mid = 0, high = nums.size() - 1;
 
        // Process elements until mid pointer crosses high pointer
        while (mid <= high) {
            // If current element is 0, swap with low and move both pointers forward
            if (nums[mid] == 0) {
                swap(nums[mid], nums[low]);
                mid++;
                low++;
            }
            // If current element is 1, it's already in correct place → move mid forward
            else if (nums[mid] == 1) {
                mid++;
            }
            // If current element is 2, swap with high and move only high pointer backward
            else {
                swap(nums[mid], nums[high]);
                high--;
            }
        }
    }
};