Sort Colors (LeetCode 75)

The problem

You get an array nums containing only 0, 1 and 2 (red, white and blue). Sort it in place so all 0s come first, then 1s, then 2s. Don’t use the library sort, and ideally do it in a single pass.

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

Two of each colour, now grouped in order.

How to think about it

Start with brute force. A normal sort is O(n log n). Counting sort is better: count the 0s, 1s and 2s, then overwrite the array. That’s O(n), but it takes two passes, and the follow-up asks for one.

The key insight: keep three zones and one unknown zone. Maintain three pointers with this invariant:

• nums[0 .. lo-1] are all 0s,
• nums[lo .. i-1] are all 1s,
• nums[hi+1 .. end] are all 2s,
• nums[i .. hi] is still unknown.

Each step looks at nums[i] and shrinks the unknown zone by one.

How each value is handled. A 0 belongs at the end of the 0-zone, so swap it with nums[lo] and move both lo and i forward (the value coming back from lo is a 1, already checked). A 2 belongs at the start of the 2-zone, so swap it with nums[hi] and move hi back, but don’t move i: the value you just pulled in from the far end hasn’t been looked at yet. A 1 is already in the right zone, so just move i.

Which pattern is this? Three-way partitioning with multiple pointers, Dijkstra’s Dutch National Flag. The tell is “group a small fixed set of values in place, in one pass”.

Approach

1. Set lo = 0, i = 0, hi = n - 1.
2. While i <= hi, look at nums[i].
3. If it’s 0: swap with nums[lo], then lo += 1 and i += 1.
4. If it’s 2: swap with nums[hi], then hi -= 1 (leave i where it is).
5. If it’s 1: i += 1.
6. When i passes hi, the unknown zone is empty and the array is sorted.

Solution

def sort_colors(nums):
    lo, i, hi = 0, 0, len(nums) - 1
    while i <= hi:                          # nums[i..hi] is still unknown
        if nums[i] == 0:
            nums[lo], nums[i] = nums[i], nums[lo]; lo += 1; i += 1
        elif nums[i] == 2:
            nums[hi], nums[i] = nums[i], nums[hi]; hi -= 1   # recheck new nums[i]
        else:
            i += 1
    return nums
# sort_colors([2,0,2,1,1,0]) -> [0,0,1,1,2,2]

Dry run

On nums = [2, 0, 2, 1, 1, 0]:

Now i > hi, so the loop ends with [0, 0, 1, 1, 2, 2].

Complexity

• Time: O(n). Every iteration either moves i forward or hi back, so the unknown zone shrinks by one each time.
• Space: O(1). Three indices; all swaps are in place.

Edge cases and pitfalls

• Empty or single element: the loop runs zero or one time.
• Only one colour ([2, 2, 2]): hi walks down to -1 with self-swaps; the array is unchanged.
• Advancing i after swapping a 2 is the classic bug. The value you swapped in could be a 0 or a 2 and still needs placing.
• Loop condition i < hi misses the last unknown element. It must be i <= hi.
• Why i can advance after swapping a 0: the value coming from lo is either a 1 (from the 1-zone) or the same element when lo == i, never an unchecked value.

Related problems

• Move Zeroes (LeetCode 283): the two-group version of in-place partitioning.
• Remove Duplicates from Sorted Array (LeetCode 26): another read/write pointer pass over an array.
• Kth Largest Element in an Array (LeetCode 215): quickselect, which uses this kind of partition step at its core.