Leetcode 75. Sort Colors

Problem

Leetcode 75 - Sort Colors

Example:

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

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

Approach

Since we know that there are only 3 colors for elements in the array, we can simply use counting sort to sort the array. We can count the frequency of each colors then we can organize the array by rearranging for each element of frequency again.

Visualization of the Approach:

nums = [2, 0, 2, 1, 1, 0]
count = [0, 0, 0] - each element represents count of 0, 1, 2

Since there are 2 0's, 1 2's and 2 2's
count = [2, 2, 2]

Now we rearrange the nums array
nums = [0, 0, 1, 1, 2, 2]

Here is the Python code for the solution:

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        #initializing the count array
        count = [0] * 3

        for n in nums:
            count[n] += 1
        
        index = 0
        # Since nums[i] is either 0, 1, 2
        for i in range(3):
            # while count[i] is more than 0, rearrange the nums array
            while count[i]:
                count[i] -= 1
                nums[index] = i
                index += 1

Time Complexity and Space Complexity

Time Complexity: $O(n)$

Space Complexity: $O(1)$ since we’re storing the space of $3$ called count[3]