We count 0’s, 1’s, and 2’s. Then fill up nums starting with 0’s, then 1’s followed by 2’s.
Time: 
class Solution:
def sortColors(self, nums: List[int]) -> None:
"""
Do not return anything, modify nums in-place instead.
"""
counts = [0, 0, 0]
for x in nums:
counts[x] += 1
curr_val = 0
for i in range(len(nums)):
while counts[curr_val] == 0:
curr_val += 1
nums[i] = curr_val
counts[curr_val] -= 1
We could sort an array around a pivot in one pass. We shall use pivot == 1 for this problem. We shall maintain nums like below:
Initially, the entire nums consists of the unknown “?” region. On each iteration, color creeps in shrinking the “?” region either from the left or from the right.
Time:
class Solution:
def sort_around_pivot(self, nums: List[int], pivot: int) -> None:
i, j, k = -1, 0, len(nums)
while j < k:
if nums[j] == pivot:
j += 1
continue
if nums[j] < pivot:
i += 1
nums[i], nums[j] = nums[j], nums[i]
j += 1
else:
k -= 1
nums[j], nums[k] = nums[k], nums[j]
def sortColors(self, nums: List[int]) -> None:
"""
Do not return anything, modify nums in-place instead.
"""
self.sort_around_pivot(nums, 1)
unreasonably effective 
Leave a comment