Sort an array of 0's 1's and 2's

Intuition

The easiest way is to sort the array using any optimal sorting algorithm. It will ensure to sort the array of 0's, 1's and 2's in ascending order.

Approach

Sort the array using built-in function for sorting. It will manage to sort the array in optimal time.

#include<bits/stdc++.h>
using namespace std;

class Solution {
public:
    //Function to sort the array
    void sortZeroOneTwo(vector<int>& nums) {
        // Sort the vector using std::sort
        sort(nums.begin(), nums.end());
    }
};
int main() {
    vector<int> nums = {2, 0, 1, 1, 0, 2};
    
    //Create an instance of Solution class
    Solution sol;
    
    sol.sortZeroOneTwo(nums);
    
    //print the array elements
    for (int num : nums) {
        cout << num << " ";
    }
    cout << endl;
    
    return 0;
}
import java.util.*;

class Solution {
    // Function to sort the array
    public void sortZeroOneTwo(int[] nums) {
        // Sort the array using Arrays.sort
        Arrays.sort(nums);
    }
}

public class Main {
    public static void main(String[] args) {
        int[] nums = {2, 0, 1, 1, 0, 2};
        
        // Create an instance of Solution class
        Solution sol = new Solution();
        
        sol.sortZeroOneTwo(nums);
        
        // Print the array elements
        for (int num : nums) {
            System.out.print(num + " ");
        }
        System.out.println();
    }
}
from typing import List
class Solution:
    # Function to sort the array
    def sortZeroOneTwo(self, nums):
        # Sort the list using sorted() function
        nums.sort()

# Main function
if __name__ == "__main__":
    nums = [2, 0, 1, 1, 0, 2]
    
    # Create an instance of Solution class
    sol = Solution()
    
    sol.sortZeroOneTwo(nums)
    
    # Print the array elements
    print(" ".join(map(str, nums)))
class Solution {
    // Function to sort the array
    sortZeroOneTwo(nums) {
        // Sort the array using Array.prototype.sort()
        nums.sort((a, b) => a - b);
    }
}

// Main function
let nums = [2, 0, 1, 1, 0, 2];

// Create an instance of Solution class
let sol = new Solution();

sol.sortZeroOneTwo(nums);

// Print the array elements
console.log(nums.join(" "));

Complexity Analysis

Time Complexity: O(NxlogN), where N is the size of the array. As the optimal sorting take O(N * logN) time. Space Complexity: O(1) no extra space is used to solve the problem.

Intuition

The better way is to keep the count of 0's, 1's and 2's. Since there are only 3 distinct values in the array so it's easy to maintain the count of all. Then we can overwrite the array based on the frequencies of 0's, 1's, 2's.

Approach

Initialize 3 variables to maintain the count of 0, 1 and 2 & iterate in the array to store the frequencies of the 0's, 1's, 2's in their corresponding variables.

Perform a 2nd traversal of array, to overwrite the array according to the frequencies, first filling up the indices by 0's, then by 1's and at last by 2's. Finally return the array.

#include<bits/stdc++.h>
using namespace std;

class Solution {
public:
    // Function to sort the array containing only 0s, 1s, and 2s
    void sortZeroOneTwo(vector<int>& nums) {
        int cnt0 = 0, cnt1 = 0, cnt2 = 0;

        // Counting the number of 0s, 1s, and 2s in the array
        for (int i = 0; i < nums.size(); i++) {
            if (nums[i] == 0) cnt0++;
            else if (nums[i] == 1) cnt1++;
            else cnt2++;
        }

        /* Placing the elements in the
        original array based on counts*/
        //placing 0's
        for (int i = 0; i < cnt0; i++) nums[i] = 0;

        //placing 1's
        for (int i = cnt0; i < cnt0 + cnt1; i++) nums[i] = 1; 
        
        //placing 2's
        for (int i = cnt0 + cnt1; i < nums.size(); i++) nums[i] = 2;
    }
};

int main() {
    vector<int> nums = {0, 2, 1, 2, 0, 1};
    
    // Create an instance of Solution class
    Solution sol;
    
    sol.sortZeroOneTwo(nums);
    
    // Print the array elements after sorting
    cout << "After sorting:" << endl;
    for (int i = 0; i < nums.size(); i++) {
        cout << nums[i] << " ";
    }
    cout << endl;
    
    return 0;
}
import java.util.*;

class Solution {
    // Function to sort the array containing only 0s, 1s, and 2s
    public void sortZeroOneTwo(int[] nums) {
        int cnt0 = 0, cnt1 = 0, cnt2 = 0;

        // Counting the number of 0s, 1s, and 2s in the array
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == 0) cnt0++;
            else if (nums[i] == 1) cnt1++;
            else cnt2++;
        }

        /* Placing the elements in the
           original array based on counts */
        // placing 0's
        for (int i = 0; i < cnt0; i++) nums[i] = 0;

        // placing 1's
        for (int i = cnt0; i < cnt0 + cnt1; i++) nums[i] = 1; 
        
        // placing 2's
        for (int i = cnt0 + cnt1; i < nums.length; i++) nums[i] = 2;
    }
    
    public static void main(String[] args) {
        int[] nums = {0, 2, 1, 2, 0, 1};
        
        // Create an instance of Solution class
        Solution sol = new Solution();
        
        sol.sortZeroOneTwo(nums);
        
        // Print the array elements after sorting
        System.out.println("After sorting:");
        for (int i = 0; i < nums.length; i++) {
            System.out.print(nums[i] + " ");
        }
        System.out.println();
    }
}
from typing import List
class Solution:
    # Function to sort the array containing only 0s, 1s, and 2s
    def sortZeroOneTwo(self, nums):
        cnt0 = 0
        cnt1 = 0
        cnt2 = 0

        # Counting the number of 0s, 1s, and 2s in the array
        for num in nums:
            if num == 0:
                cnt0 += 1
            elif num == 1:
                cnt1 += 1
            else:
                cnt2 += 1

        """ Placing the elements in the
            original array based on counts """
        # placing 0's
        for i in range(cnt0):
            nums[i] = 0

        # placing 1's
        for i in range(cnt0, cnt0 + cnt1):
            nums[i] = 1

        # placing 2's
        for i in range(cnt0 + cnt1, len(nums)):
            nums[i] = 2

if __name__ == "__main__":
    nums = [0, 2, 1, 2, 0, 1]
    
    # Create an instance of Solution class
    sol = Solution()
    
    sol.sortZeroOneTwo(nums)
    
    # Print the array elements after sorting
    print("After sorting:")
    print(" ".join(map(str, nums)))
class Solution {
    // Function to sort the array containing only 0s, 1s, and 2s
    sortZeroOneTwo(nums) {
        let cnt0 = 0;
        let cnt1 = 0;
        let cnt2 = 0;

        // Counting the number of 0s, 1s, and 2s in the array
        for (let i = 0; i < nums.length; i++) {
            if (nums[i] === 0) {
                cnt0++;
            } else if (nums[i] === 1) {
                cnt1++;
            } else {
                cnt2++;
            }
        }

        /* Placing the elements in the
           original array based on counts */
        // placing 0's
        for (let i = 0; i < cnt0; i++) {
            nums[i] = 0;
        }

        // placing 1's
        for (let i = cnt0; i < cnt0 + cnt1; i++) {
            nums[i] = 1;
        }

        // placing 2's
        for (let i = cnt0 + cnt1; i < nums.length; i++) {
            nums[i] = 2;
        }
    }
}

// Main function
let nums = [0, 2, 1, 2, 0, 1];

// Create an instance of Solution class
let sol = new Solution();

sol.sortZeroOneTwo(nums);

// Print the array elements after sorting
console.log("After sorting:");
console.log(nums.join(" "));

Complexity Analysis

Time Complexity: O(N)+O(N) = O(2N), where N is the size of the array. There are 2 traversals in the array to count the frequencies then in second iteration we are overwriting. Space Complexity: O(1) no extra space used.

Intuition

The optimal solution is a variation of the popular Dutch National flag algorithm.

This algorithm contains 3 pointers i.e. low, mid, and high, and 3 main rules.

Index 0 to low -1 contains 0
Index low to mid - 1 contains 1
Index high +1 to sizeOfArray - 1 contains 2.

The middle part i.e. mid to high is the unsorted segment. So, this part is a mix of 0's, 1's and 2's. Follow the rules mentioned in approach and image below and sort the array.

Approach

Initialize low and mid at 0 and high is sizeOfArray - 1, & iterate in the array until mid <= high. There can be three different values of mid pointer i.e. arr[mid].
- If arr[mid] is equal to 0, swap arr[low] and arr[mid] and increment both low and mid. Now the subarray from index 0 to (low-1) only contains 0.
- If arr[mid] is equal to 1, just increment the mid pointer and then the index (mid-1) will point to 1 as it should according to the rules.
- If arr[mid] is equal to 2, swap arr[mid] and arr[high] and decrement high. Now the subarray from index high+1 to (n-1) only contains 2. In this step, do nothing to the mid-pointer as even after swapping, the subarray from mid to high(after decrementing high) might be unsorted. So, check the value of mid again in the next iteration. Finally, the array should be sorted.

Dry run

For a complete dry run, refer to the video above.

#include<bits/stdc++.h>
using namespace std;

class Solution {
public:
    // Function to sort the array containing only 0s, 1s, and 2s
    void sortZeroOneTwo(vector<int>& nums) {
        
        // 3 pointers: low, mid, high
        int low = 0, mid = 0, high = nums.size() - 1; 
        while (mid <= high) {
            if (nums[mid] == 0) {
                
                /* Swap nums[low] and nums[mid], then 
                 move both low and mid pointers forward*/
                swap(nums[low], nums[mid]);
                low++;
                mid++;
                
            }
            else if (nums[mid] == 1) {
                // Move mid pointer forward
                mid++;
            }
            else {
                /* Swap nums[mid] and nums[high], 
                then move high pointer backward*/
                swap(nums[mid], nums[high]);
                high--;
            }
        }
    }
};

int main() {
    vector<int> nums = {0, 2, 1, 2, 0, 1};
    
    // Create an instance of Solution class
    Solution sol;

    sol.sortZeroOneTwo(nums);
    
    // Print the array elements after sorting
    cout << "After sorting:" << endl;
    for (int i = 0; i < nums.size(); i++) {
        cout << nums[i] << " ";
    }
    cout << endl;
    
    return 0;
}
import java.util.*;

class Solution {
    // Function to sort the array containing only 0s, 1s, and 2s
    public void sortZeroOneTwo(int[] nums) {
        // 3 pointers: low, mid, high
        int low = 0, mid = 0, high = nums.length - 1;
        
        while (mid <= high) {
            if (nums[mid] == 0) {
                
                /* Swap nums[low] and nums[mid], then 
                move both low and mid pointers forward*/
                int temp = nums[low];
                nums[low] = nums[mid];
                nums[mid] = temp;
                low++;
                mid++;
                
            } else if (nums[mid] == 1) {
                // Move mid pointer forward
                mid++;
            } else {
                /* Swap nums[mid] and nums[high], 
                then move high pointer backward*/
                int temp = nums[mid];
                nums[mid] = nums[high];
                nums[high] = temp;
                high--;
            }
        }
    }
    
    public static void main(String[] args) {
        int[] nums = {0, 2, 1, 2, 0, 1};
        
        // Create an instance of Solution class
        Solution sol = new Solution();
        
        sol.sortZeroOneTwo(nums);
        
        // Print the array elements after sorting
        System.out.println("After sorting:");
        for (int num : nums) {
            System.out.print(num + " ");
        }
        System.out.println();
    }
}
from typing import List
class Solution:
    # Function to sort the array containing only 0s, 1s, and 2s
    def sortZeroOneTwo(self, nums):
        # 3 pointers: low, mid, high
        low, mid, high = 0, 0, len(nums) - 1
        
        while mid <= high:
            if nums[mid] == 0:
                
                """ Swap nums[low] and nums[mid], then
                move both low and mid pointers forward"""
                nums[low], nums[mid] = nums[mid], nums[low]
                low += 1
                mid += 1
                
            elif nums[mid] == 1:
                
                # Move mid pointer forward
                mid += 1
                
            else:
                
                """ Swap nums[mid] and nums[high], 
                then move high pointer backward"""
                nums[mid], nums[high] = nums[high], nums[mid]
                high -= 1

# Main function for testing
if __name__ == "__main__":
    nums = [0, 2, 1, 2, 0, 1]
    
    # Create an instance of Solution class
    sol = Solution()

    sol.sortZeroOneTwo(nums)
    
    # Print the array elements after sorting
    print("After sorting:")
    print(nums)
class Solution {
    // Function to sort the array containing only 0s, 1s, and 2s
    sortZeroOneTwo(nums) {
        // 3 pointers: low, mid, high
        let low = 0, mid = 0, high = nums.length - 1;
        
        while (mid <= high) {
            if (nums[mid] === 0) {
                
                /* Swap nums[low] and nums[mid], then
                move both low and mid pointers forward*/
                [nums[low], nums[mid]] = [nums[mid], nums[low]];
                low++;
                mid++;
                
            } else if (nums[mid] === 1) {
                
                // Move mid pointer forward
                mid++;
                
            } else {
                
                /* Swap nums[mid] and nums[high], 
                then move high pointer backward*/
                [nums[mid], nums[high]] = [nums[high], nums[mid]];
                high--;
                
            }
        }
    }
}

// Main function
let nums = [0, 2, 1, 2, 0, 1];

// Create an instance of Solution class
let sol = new Solution();

sol.sortZeroOneTwo(nums);

// Print the array elements after sorting
console.log("After sorting:");
console.log(nums.join(" "));

Complexity Analysis

Time Complexity: O(N), where N is the size of the array, as there is single traversal of the array. Space Complexity: O(1), no extra space is used.

Problem List

Sort an array of 0's 1's and 2's

Examples:

Constraints

Hints

Company Tags

Editorial

Intuition

Approach

Complexity Analysis

Intuition

Approach

Complexity Analysis

Intuition

Approach

Dry run

Complexity Analysis

Frequently Occurring Doubts

Interview Followup Questions

Notes

Code