Sort Colors – Solution & Complexity

Solution Walkthrough

1. Understand the Goal

  • Sort an array containing only 0s, 1s, and 2s in a single in-place pass, without calling a general sort and without a two-pass counting sort (the follow-up explicitly asks for one pass).
  • "In-place" plus "one pass" rules out building a new array or scanning the array twice — the pointers have to make a final decision about each value as they go.

2. Choose the Core Pattern

  • The Dutch National Flag partition: maintain three pointers low, mid, high that divide the array into four regions at all times: [0, low) is all 0s, [low, mid) is all 1s, [mid, high] is unknown/unprocessed, and (high, n) is all 2s.
  • mid is the scanning cursor; low and high are the boundaries of the two settled regions, and they only ever move inward.

3. Build the Algorithm

  • While mid <= high, look at nums[mid]:
    • If it's 0: swap it with nums[low], then advance both low and mid. Advancing mid is safe: if low < mid, the value moved from low is a previously classified 1; if low == mid, the swap is a no-op. In either case, the value now at mid needs no further inspection.
    • If it's 1: it's already in the right region, just advance mid.
    • If it's 2: swap it with nums[high], then decrement high only — do not advance mid. The value swapped in from high came from the unprocessed region and hasn't been classified yet, so mid must look at it again next iteration.

4. Check Edge Cases

  • Empty or single-element array: the loop condition mid <= high is false or true-once, both handled without special-casing.
  • Already sorted ([0,0,1,1,2,2]): every comparison takes the "already in place" branch (0-swap-with-self or the 1 case), no wasted work.
  • All one color ([1,1,1]): mid just walks to the end, low/high never move.
  • All same color at the boundary, e.g. all 2s ([2,2,2]): each step swaps nums[mid] with nums[high] (sometimes a no-op self-swap) and shrinks high, until mid > high.

5. Final Solution and Complexity

  • Time complexity is O(n).
  • Space complexity is O(1).
def sort_colors(nums: list[int]) -> list[int]:
    low = 0
    mid = 0
    high = 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 -= 1
    return nums

FAQ