Three Way Flag Sort
Problem
A list holds only the values 0, 1 and 2. Arrange it so all the zeros come first, then all the ones, then all the twos. Do it inside the same list with a single pass over the values, without counting how many of each value there are first.
Examples
Input: values = [2, 0, 2, 1, 1, 0]
Output: [0, 0, 1, 1, 2, 2]Input: values = [2, 0, 1]
Output: [0, 1, 2]
Why: every value is a different one, so each has to moveInput: values = []
Output: []
Why: edge case, there is nothing to arrangeHints
0 / 3
Only three distinct values exist, so a general comparison sort is far more machinery than the problem needs.
Think of the list as four regions that grow and shrink: the finished zeros, the finished ones, the part not yet looked at, and the finished twos at the far end. Three moving boundaries describe them.
Inspect the first unexamined value. A zero is swapped down to the boundary of the zeros and both that boundary and the inspection point advance. A two is swapped up to the boundary of the twos, which moves inward while the inspection point stays put because the value swapped in is still unexamined. A one needs no swap, so the inspection point simply advances.
Solution
Three boundaries carve the list into finished zeros, finished ones, unexamined values, and finished twos. Each step inspects one unexamined value and places it in constant time: zeros swap down to the low boundary, twos swap up to the high boundary, and ones are already where they belong. The asymmetry matters — after swapping a two into place the inspection point must not advance, because the value it received from the far end has never been looked at. Time is O(n) in a single pass, and space is O(1).
def flag_sort(values):
low, mid, high = 0, 0, len(values) - 1
while mid <= high:
if values[mid] == 0:
values[low], values[mid] = values[mid], values[low]
low, mid = low + 1, mid + 1
elif values[mid] == 2:
values[mid], values[high] = values[high], values[mid]
high -= 1 # the value swapped in is still unexamined
else:
mid += 1 # a one is already in the middle region
return values
print(flag_sort([2, 0, 2, 1, 1, 0])) # -> [0, 0, 1, 1, 2, 2]
print(flag_sort([2, 0, 1])) # -> [0, 1, 2]
print(flag_sort([])) # -> []Stuck on the idea rather than the code? Dutch National Flag covers it.