Sort Colors

medium array three pointers dutch flag

Problem

Given an array nums with n objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue. We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.

Inputnums = [2, 0, 2, 1, 1, 0]

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

Three pointers low, mid, high partition the array into three regions.

nums (only 0, 1, 2 — comma-separated)

def sort_colors(nums):
    low, mid, high = 0, 0, 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] == 2:
            nums[mid], nums[high] = nums[high], nums[mid]
            high -= 1
        else:
            mid += 1
    return nums

function sortColors(nums) {
  let low = 0, mid = 0, high = nums.length - 1;
  while (mid <= high) {
    if (nums[mid] === 0) {
      [nums[low], nums[mid]] = [nums[mid], nums[low]];
      low++; mid++;
    } else if (nums[mid] === 2) {
      [nums[mid], nums[high]] = [nums[high], nums[mid]];
      high--;
    } else {
      mid++;
    }
  }
  return nums;
}

class Solution {
    public void sortColors(int[] nums) {
        int low = 0, mid = 0, high = nums.length - 1;
        while (mid <= high) {
            if (nums[mid] == 0) {
                int t = nums[low]; nums[low] = nums[mid]; nums[mid] = t;
                low++; mid++;
            } else if (nums[mid] == 2) {
                int t = nums[mid]; nums[mid] = nums[high]; nums[high] = t;
                high--;
            } else {
                mid++;
            }
        }
    }
}

void sortColors(vector<int>& nums) {
    int low = 0, mid = 0, high = (int)nums.size() - 1;
    while (mid <= high) {
        if (nums[mid] == 0) {
            swap(nums[low], nums[mid]);
            low++; mid++;
        } else if (nums[mid] == 2) {
            swap(nums[mid], nums[high]);
            high--;
        } else {
            mid++;
        }
    }
}

Explanation

We only have three values (0, 1, 2), so we can sort them in a single pass with no counting and no extra memory. This is the Dutch National Flag technique.

We keep three pointers. Everything before low is a confirmed 0, everything after high is a confirmed 2, and mid is the element we are currently inspecting. We loop while mid <= high.

At each step we look at nums[mid]. If it is 0, we swap it down to the low region and advance both low and mid. If it is 2, we swap it up to the high region and shrink high — but we do not move mid, because the value we just pulled in still needs checking. If it is 1, it already belongs in the middle, so we just advance mid.

This works because the three pointers always keep the array split into "done 0s | unknown | done 2s", and that unknown gap shrinks every iteration until it disappears.

Example: [2,0,2,1,1,0]. The first 2 swaps to the back, a 0 swaps to the front, and so on, ending as [0,0,1,1,2,2] after one sweep.

Time: O(n) Space: O(1)