Sort an array of 0s, 1s and 2s in one pass: three pointers sweep zeros to the front and twos to the back, leaving ones in the middle.
▼The problem
LeetCode 75 (Medium). Given an array nums of n values, each 0, 1 or 2 (red, white and blue), sort it in place so that equal colours are adjacent, in the order 0, 1, 2, without using the library's sort. The follow-up asks for a one-pass algorithm with constant extra space.
Examples (LeetCode's): nums = [2, 0, 2, 1, 1, 0] → [0, 0, 1, 1, 2, 2] (walked through in scene 6) and nums = [2, 0, 1] → [0, 1, 2].
The solution
def sortColors(nums):
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 -= 1Transcript
Sort Colors. A hen house row holds eggs marked zero, one or two: red, white and blue. Sort them in place, so equal colours sit together in that order, without a library sort.
Take two, zero, two, one, one, zero. Sorted, it becomes zero, zero, one, one, two, two. And two, zero, one becomes zero, one, two.
Counting works: two zeros, two ones, two twos, then rewrite the row. Linear, but two passes. A comparison sort is n log n. The follow-up wants one pass and constant space.
The trick is Dijkstra's Dutch national flag, with three pointers. Before low is all zeros, after high all twos, and from low up to mid, ones. Mid checks the next unknown egg. A zero swaps to low, and both move up. A one just lets mid move on. A two swaps to high, and high moves down, but mid stays: the egg it got back is unchecked.
In code, it's one loop that runs while mid is at most high, with three branches.
Let's run it. Mid finds a two and swaps it with the zero at the end; high steps in. Mid checks that zero: it swaps with low, and both advance. The next zero does the same. Then a two swaps with the one before high. Mid walks past two ones and passes high. Done: zero, zero, one, one, two, two.
Every step moves mid up or high down, so it's one pass: order n time, and order one space, just three pointers.
Zeros to the front, twos to the back, ones in the middle. That's Sort Colors.