Sort Colors

Intuition

Picture the array as three growing zones: settled zeros on the left, settled twos on the right, and an unexplored middle. A single scanning pointer looks at each value and either drops it into the zeros zone by swapping left, leaves it in place as a one, or throws it into the twos zone by swapping right. Because a value swapped in from the right has not been examined yet, we re-inspect that position, but every 0 or 2 we place is final, so the whole thing finishes in one linear pass.

Algorithm walkthrough

1. 1Keep three indices: low marks the next slot for a 0, high marks the next slot for a 2, and mid is the current position under inspection.
2. 2While mid has not passed high, read nums[mid].
3. 3If it is 0, swap it with nums[low] and advance both low and mid, since the swapped-in value was already checked.
4. 4If it is 1, leave it in place and just advance mid.
5. 5If it is 2, swap it with nums[high] and shrink high, but do not move mid because the incoming value still needs a look.

Implementation

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        low, mid, high = 0, 0, len(nums) - 1
        while mid <= high:
            if nums[mid] == 0:
                nums[low], nums[mid] = nums[mid], nums[low]
                low += 1
                mid += 1
            elif nums[mid] == 1:
                mid += 1
            else:
                nums[mid], nums[high] = nums[high], nums[mid]
                high -= 1

Intuition

Since there are only three distinct values, comparisons are unnecessary. One pass counts how many 0s, 1s, and 2s appear. A second pass then overwrites the array from left to right: first as many 0s as we counted, then the 1s, then the 2s. The array ends up sorted because the colors are emitted in ascending order.

Algorithm walkthrough

1. 1Create a size-three tally array initialized to zero.
2. 2Scan the input once, incrementing the tally for each value seen.
3. 3Start an output index at the front of the array.
4. 4For each color 0, then 1, then 2, write that color into the array as many times as its tally, advancing the output index each time.

Implementation

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        counts = [0, 0, 0]
        for value in nums:
            counts[value] += 1
        idx = 0
        for color in range(3):
            for _ in range(counts[color]):
                nums[idx] = color
                idx += 1