Sort Colors

Problem

Write a function to sort a given integer array nums in-place (and without the built-in sort function), where the array contains n integers that are either 0, 1, and 2 and represent the colors red, white, and blue. Arrange the objects so that same-colored ones are adjacent, in the order of red, white, and blue (0, 1, 2).

Input:

nums = [2,1,2,0,1,0,1,0,1]

Output:

[0,0,0,1,1,1,1,2,2]

Explanation

We can understand this algorithm by looking at the invariants which hold true after each iteration:

• All elements to the left of the left are 0s.
• All elements between left and i - 1 are 1s.
• All elements between i and right are unsorted.
• All elements to the right of right are 2s.

Let’s now see how we maintain these invariants as we iterate through the array.

Sorting 0s

When nums[i] is equal to 0, invariant 2 tells us there are two possible values for left: 0 or 1.

Let’s consider the case when left == 1 first. We swap i with left. This allows us to increment the left pointer to maintain invariant 1. Since we know that the new item at i is a 1, we can also increment i to maintain variant 2.

Now let’s consider what happens when left == 0, which happens when we haven’t encountered any 1s yet and i == left. We swap i with left (which is itself) and increment left to maintain invariant 1. Since we still haven’t encountered any 1s, we can increment i to maintain invariant 2. In other words, the “ones” region remains empty.

Sorting 1s

When nums[i] == 1, we can simply increment i to maintain invariant 2.

Sorting 2s

When nums[i] == 2, we swap nums[i] with nums[right]. This allows us to decrement right to maintain invariant 4. But since the new item at i came from the unsorted region, the new item at i is still unsorted, so we have to go through another iteration to correctly sort it.

Termination

When i surpasses right the unsorted region is empty and the entire array is sorted.

Solution

def sortColors(nums):
  left, right = 0, len(nums) - 1
  i = 0

  while i <= right:
    if nums[i] == 0:
      nums[i], nums[left] = nums[left], nums[i]
      left += 1
      i += 1
    elif nums[i] == 2:
      nums[i], nums[right] = nums[right], nums[i]
      right -= 1
    else:
      i += 1

public void sortColors(int[] nums) {
    int left = 0, right = nums.length - 1;
    int i = 0;

    while (i <= right) {
        if (nums[i] == 0) {
            int temp = nums[i];
            nums[i] = nums[left];
            nums[left] = temp;
            left++;
            i++;
        } else if (nums[i] == 2) {
            int temp = nums[i];
            nums[i] = nums[right];
            nums[right] = temp;
            right--;
        } else {
            i++;
        }
    }
}

func sortColors(nums []int) {
    left, right := 0, len(nums)-1
    i := 0

    for i <= right {
        if nums[i] == 0 {
            nums[i], nums[left] = nums[left], nums[i]
            left++
            i++
        } else if nums[i] == 2 {
            nums[i], nums[right] = nums[right], nums[i]
            right--
        } else {
            i++
        }
    }
}

function sortColors(nums: number[]): void {
    let left = 0, right = nums.length - 1;
    let i = 0;

    while (i <= right) {
        if (nums[i] === 0) {
            [nums[i], nums[left]] = [nums[left], nums[i]];
            left++;
            i++;
        } else if (nums[i] === 2) {
            [nums[i], nums[right]] = [nums[right], nums[i]];
            right--;
        } else {
            i++;
        }
    }
}