Sort Colors

Solution

Two Pointers

This problem is basically the same as Dutch National Flag

Wrote a LeetCode post (opens new window).

def sortColors(self, nums: List[int]) -> None:
    """
    Do not return anything, modify nums in-place instead.
    """
    left, right = 0, len(nums)-1
    # inv: nums[left..t-1] are 0's, nums[t..i-1] are undecided, nums[i..j] are 1's, nums[j+1..right] are 2's
    t = left
    i = right + 1
    j = right

    while t < i:
        if nums[i-1] == 0:
            nums[i-1], nums[t] = nums[t], nums[i-1]
            t += 1
        elif nums[i-1] == 1:
            i -= 1
        else:
            nums[i-1], nums[j] = nums[j], nums[i-1]
            i -= 1
            j -= 1
    # post: nums[left..i-1] are 0's, nums[i..j] are 1's, nums[j+1..right] are 2's