# 75. Sort Colors

## Problem

• Given an array 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.

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

• Note: You are not suppose to use the library’s sort function for this problem.

Example:

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


## Thought Process

• Loop through the whole array
• count how many 0s, 1s and 2s are there
• reorder the array as 0s + 1s + 2s

## White Board

Below is the white board:

## Code

class Solution:
def sortColors(self, nums):
color_dict = {}
for num in nums:
assert num in (0, 1, 2) #make sure num is valid
if num in color_dict:
color_dict[num] += 1
else:
color_dict[num] = 1

idx = 0
for i in range(color_dict[0]):
nums[idx] = 0
idx += 1
for i in range(color_dict[1]):
nums[idx] = 1
idx += 1
for i in range(color_dict[2]):
nums[idx] = 2
idx += 1

def test_sortColors():
nums = [0,1,2,2,1,1,2,2,0,0,0,0,2,1]
print("Before Sort: ")
print(nums)
sc = Solution()
sc.sortColors(nums)
print("After Sort: ")
print(nums)

test_sortColors()

Before Sort:
[0, 1, 2, 2, 1, 1, 2, 2, 0, 0, 0, 0, 2, 1]
After Sort:
[0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2]


## Complexity

• Time complexity is O(n), as a single for loop iteration is used
• Space complexity is O(k), k is the number of different colors, as a dictionary of size k was created

## Optimization

Here the 2nd half of the solution is tedious to loop thru each color. An alternative is to use a 3-ways Quick Sort. Below is an illustration how it works:

For a given state:

• There are three segments:

• nums[0, lt - 1], smaller than 1 (0s)
• nums[lt, i - 1], all 1s
• nums[gt, n-1], all 2s
• i is the current position, current value is nums[i]

• if nums[i] == 1, do nothing, i ++
• if nums[i] == 0, swap with the first 1, a.k.a nums[lt], update lt++, i++
• if nums[i] == 2, swap with the value before first 2, a.k.a nums[gt-1], update gt–, no need to update i as we swap the value before first 2, which is not evaluated at this moment
class Solution2:

def sortColors2(self, nums):
lt = 0
# zeros = nums[0, lt - 1], initial as -2 so the initial segment is nums[0, -1] which is invalid
gt = len(nums)
# twos = nums[gt, n-1], initial as n so the initial segment is nums[n, n-1] which is invalid
idx_i = 0
for i in range(gt):
if nums[idx_i] == 1:
idx_i += 1
elif nums[idx_i] < 1:
# if nums[idx] == 0
assert nums[idx_i] == 0
# make sure the input value is valid
nums[idx_i], nums[lt - 1] = nums[lt - 1], nums[idx_i]
lt += 1
idx_i += 1
else:
assert nums[idx_i] == 2
# make sure the input value is valid
nums[idx_i], nums[gt - 1] = nums[gt - 1], nums[idx_i]
gt -= 1
# note no need to update idx_i here

def test_sortColors2():
nums = [0,1,2,2,1,1,0,2,0,0,0,0,2,1,0,2]
print("Before Sort: ")
print(nums)
sc = Solution2()
sc.sortColors2(nums)
print("After Sort: ")
print(nums)

test_sortColors2()

Before Sort:
[0, 1, 2, 2, 1, 1, 0, 2, 0, 0, 0, 0, 2, 1, 0, 2]
After Sort:
[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2]


Time complexity is O(n) and space complexity is O(1), but there are much fewer operations