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