75. Sort Colors (Python)

Related Topic

Array. Two-Pointers.

Description

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.

Sample I/O

Example 1

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

Follow up:

A rather straight forward solution is a two-pass algorithm using counting sort. First, iterate the array counting number of 0’s, 1’s, and 2’s, then overwrite array with total number of 0’s, then 1’s and followed by 2’s. Could you come up with a one-pass algorithm using only constant space?

Methodology

The hardest part of this question is these constrains. Have to be certain order (0,1,2), in-place, one-pass and no build-in sorting library. Without these constrains, this question is pretty easy, you can:

• Return directly sorted list
• Use dictionary to count each element and overwrite the original array

The standard solution is associated with dutch national flag problem. We know the leftmost must be smallest among three numbers (here is 0), the rightmost must be largest (here is 2) and the rest can only be 1. We need three pointers to indicate the left boundary, current and right boundary. We iterate the list, if the current == 0, swap current with left boundary and update left boundary and current pointers; if the current == 1, we update the current pointer; if current == 2, we swap the current and right boundary and update right boundary pointers.

def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        left = cur = 0
        right = len(nums)-1
        while cur<=right:
            if nums[cur]==0:
                nums[left],nums[cur]=nums[cur],nums[left]
                left+=1
                cur+=1
            elif nums[cur]==2:
                nums[cur],nums[right]=nums[right],nums[cur]
                right-=1
            else:
                cur+=1

BigO

Time complexity : O(nlogn) since it’s one pass along the array of length n.