# 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