Topics › Arrays › Medium

Sort an array of 0's 1's and 2's

Medium TUF ID 36 SDE

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

You are given an array nums whose values are only 0, 1 and 2. Sort it in place so that all 0s come first, then all 1s, then all 2s. The intended solution does not call a library sort. (On LeetCode this is “Sort Colors”: red, white and blue are encoded as 0, 1, 2.)

Constraints: 1 ≤ n ≤ 300 on LeetCode, but the techniques below are meant for arrays of any size, and the Java file also handles the empty array.

Example 1

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

Example 2

nums = [2, 0, 1]
after = [0, 1, 2]        // each value appears once

The follow-up question is the interesting part: can you do it in one pass with constant extra space?

Approach 1: Brute force, use a general sort

Intuition. It is a sorting problem, so any comparison sort works. This ignores the most useful fact we have, that there are only three distinct values, and so it pays the general O(n log n) price.

Algorithm.

1. Call Arrays.sort(nums).

static void bruteForce(int[] nums) {
    Arrays.sort(nums);
}

Complexity. Time O(n log n). Space O(log n) for the sort’s recursion on primitive arrays.

In an interview this is a starting point to mention, not an answer.

Approach 2: Better, counting

Intuition. With only three possible values, the sorted array is fully described by three numbers: how many 0s, 1s and 2s there are. Count them in one pass, then overwrite the array: the first cnt[0] slots with 0, the next cnt[1] with 1, and the rest with 2. This is counting sort with a three-bucket table.

Algorithm.

1. Count occurrences of each value into cnt[0..2].
2. Walk v = 0, 1, 2 and write v exactly cnt[v] times, advancing a write index k.

static void better(int[] nums) {
    int[] cnt = new int[3];
    for (int v : nums) cnt[v]++;
    int k = 0;
    for (int v = 0; v < 3; v++) {
        for (int c = 0; c < cnt[v]; c++) nums[k++] = v;
    }
}

Complexity. Time O(n), two passes over the array. Space O(1), a fixed array of three counters.

This is already linear and is perfectly acceptable. The only thing it does not achieve is the one-pass follow-up, and it rewrites values rather than moving the original elements, which matters if the elements carry other data (for example, objects keyed by colour).

Approach 3: Optimal, Dutch National Flag

Intuition. Edsger Dijkstra’s Dutch National Flag algorithm partitions the array into three regions with three pointers, while a single scan shrinks the unknown middle part. At every moment:

[0, low)       all 0s
[low, mid)     all 1s
[mid, high]    not examined yet
(high, n)      all 2s

Look at nums[mid], the first unexamined element:

• It is 0. It belongs in the 0-region. Swap it with nums[low], which is the first 1 (or mid itself if there are no 1s yet). Both regions grow: low++, mid++. The element that came back to mid is a known 1 (or the same 0), so it is safe to advance.
• It is 1. It is already in the right place. mid++.
• It is 2. It belongs at the end. Swap it with nums[high] and do high--. Do not advance mid: the value that just arrived from high has not been examined yet.

The loop ends when the unknown region is empty, that is, when mid > high.

Algorithm.

1. low = 0, mid = 0, high = n - 1.
2. While mid ≤ high, apply the three cases above.

static void optimal(int[] nums) {
    int low = 0, mid = 0, high = nums.length - 1;
    // [0, low) are 0s, [low, mid) are 1s, [mid, high] unknown, (high, n) are 2s
    while (mid <= high) {
        if (nums[mid] == 0) {
            swap(nums, low++, mid++);
        } else if (nums[mid] == 1) {
            mid++;
        } else {
            swap(nums, mid, high--);       // do not advance mid: the swapped-in value is unexamined
        }
    }
}

static void swap(int[] a, int i, int j) {
    int t = a[i];
    a[i] = a[j];
    a[j] = t;
}

Dry run on nums = [2, 0, 2, 1, 1, 0]:

Complexity. Time O(n): every iteration either advances mid or shrinks high, so the unknown region loses one element per iteration and there are at most n iterations. Space O(1).

Why it is correct. The four-region picture is an invariant: it holds before the loop (three regions are empty and everything is unknown) and each case preserves it. When the unknown region is empty, the remaining three regions are exactly the sorted order.

Edge cases and pitfalls

• Advancing mid after swapping with high. This is the classic bug. The incoming element might be a 0 or a 2 and would be skipped. Only the 0 case and the 1 case advance mid.
• Loop condition. Use mid <= high, not mid < high; the element at high is still unknown and must be examined.
• All one value. [2, 2, 2] makes high walk down to -1 while mid stays at 0; [0, 0, 0] makes low and mid walk together. Both terminate correctly.
• Empty array. high = -1, the loop never runs.
• Values outside 0..2. The counting version would throw an index error and the flag version would treat them as 2. The problem rules them out; validate if the input is untrusted.
• Stability. The Dutch National Flag swaps elements over long distances, so records with equal keys can change their relative order. For plain integers this is invisible; for records keyed by colour it can matter.

Related problems

• Sort an array of 0’s 1’s and 2’s is the same problem as it appears in the A2Z sheet.
• Sort a Linked List of 0’s 1’s and 2’s applies the three-bucket idea to linked lists.
• Quick Sorting uses a partition step; three-way partitioning is the Dutch National Flag.
• Majority Element-I is another single-pass, constant-space array scan with a carefully maintained invariant.