Sort Colors

Q: What is the Dutch national flag problem?

Edsger Dijkstra posed it: given objects of three colors in a row, the red, white and blue of the Dutch flag, group each color together in one pass, using only swaps. Sort Colors is the same problem with the numbers 0, 1 and 2. His solution is the three-pointer partition with low , mid and high .

Q: What is the time and space complexity of Sort Colors?

The one-pass solution runs in O(n) time, because every step shrinks the unread part by one cell. It uses O(1) extra space: three indices and a temporary value for the swap. Counting sort has the same bounds but reads the array twice.

Q: Why does mid not move after swapping with high?

The value that comes back from high has never been read, so it could be a 0, a 1 or a 2. Moving mid past it would leave a 0 or a 2 in the middle. A swap with low is different: everything between low and mid is a 1, so the value that comes back is known and mid can move on.

Q: Is counting sort an acceptable answer for Sort Colors?

It meets the O(n) time and O(1) space bounds, and many interviewers accept it as a first answer. Expect the follow-up asking for a single pass, which is the three-pointer partition. Counting is the better tool when there are many colors, since the partition only splits into three groups.

MediumTwo pointers Sorting

+10

You get an array nums in which every value is 0, 1 or 2. Think of them as three colors, for example red, white and blue. Rearrange the array so that all the 0s come first, then all the 1s, then all the 2s, and return it.

Solve it without a library sort function. The point is to use what you know about the values.

Function

sortColors(nums: integer-array) → integer-array

numsinteger-array: the colors, each one 0, 1 or 2
Returnsinteger-array: the same values with every 0 first, then every 1, then every 2

Constraints

1 ≤ nums.length ≤ 1.5 × 10⁴
Every nums[i] is 0, 1 or 2.
A color may be missing, and the array may hold a single color.

Examples

Input: nums = [2, 1, 0, 2, 0, 1, 1]
Output: [0, 0, 1, 1, 1, 2, 2]
Explanation: The array holds two 0s, three 1s and two 2s, so the result is exactly that: two 0s, then three 1s, then two 2s.

+17 hidden tests on Submit

Follow-up

What would you change if there were k colors instead of three, with k much smaller than the length of the array?

Hints

Open them one at a time. Each one gives away a little more.

0 of 3 opened

Solution

Any sort gives the right order, so the real question is what the three values let you skip. Because only 0, 1 and 2 can appear, you can count them and rewrite the array in two passes. With three pointers that mark where the 0s end and where the 2s begin, you can even put every value in its place in a single pass. That one-pass partition is the Dutch national flag algorithm.

Bubble sort by hand

Correct, but does not finish on the largest tests

Intuition

A library sort would pass in O(n log n), but the problem rules it out, because an interviewer wants to see what you do with the fact that there are only three values. The baseline is then a sort you write yourself, and the shortest one to get right is bubble sort: walk the array, and whenever two neighbors are out of order, swap them.

One pass carries the largest value it meets all the way to the end, like a bubble rising. After the first pass the last position is final, after the second the last two are, so n-1 passes leave the whole array in order. In [2, 1, 0] the first pass moves the 2 to the end, giving [1, 0, 2], and the second pass swaps the 1 and the 0.

It is slow because each pass compares every pair that is not final yet: about n²/2 comparisons in total. With n = 1.5 × 10^4 that is over 10^8 comparisons, plus a swap for every pair that starts out of order, and none of that work uses the fact that only three values exist.

Algorithm

Run n-1 passes over the array.
In each pass, compare every pair of neighbors nums[j] and nums[j + 1] that is not final yet, and swap them when the left one is bigger.
After pass number done (counting from 0), the last done + 1 positions hold their final values, so the next pass stops before them.
Return nums.

def sortColors(nums):
    n = len(nums)
    for done in range(n - 1):
        # One pass: the largest value left so far bubbles to index n-1-done.
        for j in range(n - 1 - done):
            if nums[j] > nums[j + 1]:
                nums[j], nums[j + 1] = nums[j + 1], nums[j]
    return nums

Count each color, then rewrite

Intuition

Bubble sort spends all its time comparing neighbors, but you already know which values exist. If the array holds two 0s, three 1s and two 2s, the answer is fixed before you move anything: two 0s, three 1s, two 2s. Only the counts matter.

So read the array once and count each value. Then overwrite it from the start: count[0] zeros, then count[1] ones, then count[2] twos. This is counting sort, and it is safe here because equal values are interchangeable. A 1 is a 1, so nothing about the original order has to survive.

That is two passes and three counters, O(n) time and O(1) space. It meets the bounds, and it is the natural answer when there are many colors. The follow-up this problem is known for is whether you can do it while reading the array only once.

Algorithm

Create three counters, all 0.
Read every value and add one to its counter.
Write count[0] zeros from the start, then count[1] ones, then count[2] twos.
Return nums.

def sortColors(nums):
    count = [0, 0, 0]  # how many 0s, 1s and 2s
    for x in nums:
        count[x] += 1
    i = 0
    for color in range(3):
        for _ in range(count[color]):
            nums[i] = color
            i += 1
    return nums

One pass with three pointers (Dutch national flag)

Intuition

Grow three regions while you read: 0s at the front, 1s after them, 2s at the back, and an unread part between the 1s and the 2s. Three indices mark the borders. Everything before low is 0, everything from low up to but not including mid is 1, everything after high is 2, and nums[mid] to nums[high] is still unread.

Read nums[mid]. A 1 is already in its region, so move mid on. A 0 belongs at the front: swap it with nums[low] and move both low and mid on. The value that comes back from low is a 1 (or the same 0, when no 1 has been seen yet), so it is already in place. A 2 belongs at the back: swap it with nums[high] and move high back, but keep mid where it is, because the value that came from high has not been read.

Every step moves mid forward or high back, so the unread part loses one cell each time and the loop ends after n steps. Trace [2, 0, 2]: the first 2 swaps with the last 2 and high drops to 1; index 0 still holds a 2, which swaps with the 0 and high drops to 0; index 0 now holds the 0, which stays, and you get [0, 2, 2].

Algorithm

Set low = 0, mid = 0 and high to the last index.
While mid ≤ high, read nums[mid].
If it is 0, swap it with nums[low] and move low and mid one step right.
If it is 1, move mid one step right.
If it is 2, swap it with nums[high] and move high one step left. Leave mid in place.
Return nums.

def sortColors(nums):
    # nums[:low] are 0s, nums[low:mid] are 1s, nums[high + 1:] are 2s.
    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] == 1:
            mid += 1
        else:
            # The value swapped in from high is unread, so mid stays.
            nums[mid], nums[high] = nums[high], nums[mid]
            high -= 1
    return nums

Pitfalls and edge cases

The one-pass version is short, and almost every bug in it is a pointer that moves when it should not.

Moving mid forward after a swap with high. The value that arrives is unread. On [1, 2, 0] the 2 swaps with the 0, and skipping the 0 returns [1, 0, 2].
Looping while mid < high when high is the last unread index. When the two meet, that cell is still unread. On [1, 0] the loop stops before it reads the 0 and returns [1, 0].
Letting high fall below zero with an unsigned index. An array of only 2s, like [2], drives high to -1. In Rust, where indices are usize, keep high one past the unread part instead, as the Rust code does.
Assuming every color appears. [2, 0, 2] has no 1, and an array can hold a single color. The pointer rules handle both without special cases, so do not add any.

Frequently asked questions4

What is the Dutch national flag problem?

Edsger Dijkstra posed it: given objects of three colors in a row, the red, white and blue of the Dutch flag, group each color together in one pass, using only swaps. Sort Colors is the same problem with the numbers 0, 1 and 2. His solution is the three-pointer partition with low, mid and high.

What is the time and space complexity of Sort Colors?

The one-pass solution runs in O(n) time, because every step shrinks the unread part by one cell. It uses O(1) extra space: three indices and a temporary value for the swap. Counting sort has the same bounds but reads the array twice.

Why does mid not move after swapping with high?

The value that comes back from high has never been read, so it could be a 0, a 1 or a 2. Moving mid past it would leave a 0 or a 2 in the middle. A swap with low is different: everything between low and mid is a 1, so the value that comes back is known and mid can move on.

Is counting sort an acceptable answer for Sort Colors?

It meets the O(n) time and O(1) space bounds, and many interviewers accept it as a first answer. Expect the follow-up asking for a single pass, which is the three-pointer partition. Counting is the better tool when there are many colors, since the partition only splits into three groups.

Python

def sortColors(nums):
    # Write code here

Test cases

Case 1

Case 2

Case 3

Input

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

Expected

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

Sort Colors

Function

Constraints

Examples

Hints

Solution

Bubble sort by hand

Count each color, then rewrite

One pass with three pointers (Dutch national flag)

Pitfalls and edge cases

Frequently asked questions4

Similar problems