Similar principle as Bucket Sort. We use subarrays of nums as “buckets” instead of creating an array of list, because we are swapping in-place.
class Solution:
def sortColors(self, nums: List[Literal[0] | Literal[1] | Literal[2]]) -> None:
"""
Do not return anything, modify nums in-place instead.
"""
# INVARIANT: nums consists of 0, 1, 2
# NOTE:
# i: Count of elements that are 0 in sorted subarray
# j: Count of elements that are 1 in sorted subarray
# k: Count of elements that are 2 in sorted subarray
i=j=k = 0
while (curr_idx := i + j + k) < len(nums):
n = nums[curr_idx]
if n == 0:
nums[i], nums[curr_idx] = nums[curr_idx], nums[i]
i += 1
# INVARIANT: 0 is swapped with 1 xor 2 that is an element in the sorted subarray
if nums[curr_idx] == 1:
j = max(0, j-1)
else:
k = max(0, k-1)
if n == 1:
nums[i+j], nums[curr_idx] = nums[curr_idx], nums[i+j]
j += 1
# INVARIANT: 1 is swapped with 2 that is an element in the sorted subarray
k = max(0, k-1)
if n == 2:
nums[i+j+k], nums[curr_idx] = nums[curr_idx], nums[i+j+k]
k += 1