🌈75. Sort Colors

Medium | Grind 75 | February 15th Wednesday, 2023

Given an array nums with n objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue.

We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.

You must solve this problem without using the library's sort function.

Constraints:

• n == nums.length

• 1 <= n <= 300

• nums[i] is either 0, 1, or 2.

Follow up: Could you come up with a one-pass algorithm using only constant extra space?

Example 1

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

Example 2

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

SOLUTION

We set up two index variables: zero is set to 0 to keep track of the current position of the last 0 that was sorted, and two is set to the last index in the array (len(nums)-1) to keep track of the current position of the first 2 that has been sorted.

Next, we iterate through the array until the iterator is more than the value of two. As we iterate over the elements, if it is equal to 0, then the current element is swapped with the element at the current position of zero, and both the iterator and zero are incremented by 1. This moves the 0 to the left side of the array. If the current element is equal to 2, then the current element is swapped with the element at the current position of two, and two is decremented by 1. This moves the 2 to the right side of the array. If the current element is equal to 1, then it is already in the correct position, and the iterator is incremented by 1.

After the while loop finishes, the array will have all the 0s on the left, followed by all the 1s, and then all the 2s on the right.

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        
        zero = 0 # index to track 0 position
        two = len(nums)-1 # index to track 2 position

        iterator = 0

        while iterator <= two:
            # move the 0 to left 
            if nums[iterator] == 0: 
                nums[iterator], nums[zero] = nums[zero], nums[iterator]
                iterator += 1
                zero += 1
            
            # move the 2 to the right
            elif nums[iterator] == 2: 
                nums[iterator], nums[two] = nums[two], nums[iterator]
                two -= 1
                
            # we do not care about the 1
            else: 
                iterator += 1
        
        return nums

TIME AND SPACE COMPLEXITY

Time: O(n) where n is the length of the input array. This is because the algorithm makes a single pass through the array, and each element is visited at most twice (once during the swap with the 0 or 2, and once during the iterator increment).

SPACE: O(1)This is because the algorithm only uses a few additional variables (namely zero, two, and iterator) to keep track of the position of the 0s, 1s, and 2s, and no additional data structures are used to store the input array or intermediate results.

**written together with ChatGPT