Topics › Arrays › Medium

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

Medium TUF ID 2762 A2Z

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

You are given an array nums whose elements are only 0, 1 or 2. Rearrange it in place so that all 0s come first, then all 1s, then all 2s. The LeetCode version calls the values colors (red, white, blue) and the problem “Sort Colors”; the well-known one-pass solution is the Dutch National Flag algorithm, named after the three horizontal bands of that flag.

The interesting constraint is the unspoken one: do it in a single pass with constant extra memory, without calling a library sort.

Constraints: 0 ≤ n ≤ 3 × 10^5, nums[i] ∈ {0, 1, 2}.

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]

Example 3

nums = [1, 1, 1]
after  = [1, 1, 1]      // nothing moves

Approach 1: Brute force, a general sort

Intuition. The task is literally “sort the array”, so a general-purpose sort solves it. It ignores the fact that there are only three distinct values, and that is exactly what the next approaches exploit.

Algorithm.

1. Call the library sort.

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

Complexity. Time O(n log n). Space O(log n) for the recursion of the dual-pivot quicksort that Arrays.sort(int[]) uses.

Approach 2: Better, counting

Intuition. With only three possible values, sorting is the same as knowing how many of each there are. Count the zeros, ones and twos in one pass, then overwrite the array: c0 zeros, then c1 ones, then c2 twos. This is counting sort specialised to three buckets.

Algorithm.

1. Count c0, c1, c2.
2. Write 0 into the first c0 positions, 1 into the next c1, and 2 into the remaining c2.

static void better(int[] nums) {
    int c0 = 0, c1 = 0, c2 = 0;
    for (int x : nums) {
        if (x == 0) c0++;
        else if (x == 1) c1++;
        else c2++;
    }
    int i = 0;
    while (c0-- > 0) nums[i++] = 0;
    while (c1-- > 0) nums[i++] = 1;
    while (c2-- > 0) nums[i++] = 2;
}

Complexity. Time O(n): two passes over the array. Space O(1): three counters.

This is already linear and is a perfectly good interview answer. Its only shortcoming is the second pass: the data has to be read completely before any output is written. If the array is being streamed to you, or if the follow-up asks for “one pass”, you need the next idea. Note also that the write loops post-decrement the counters: while (c0-- > 0) runs exactly c0 times.

Approach 3: Optimal, Dutch National Flag

Intuition. Maintain three regions in the array using two boundaries and a cursor:

• [0, low) holds only 0s,
• [low, mid) holds only 1s,
• [mid, high] is unknown, not yet examined,
• (high, n-1] holds only 2s.

Initially low = mid = 0 and high = n - 1, so everything is unknown. Look at nums[mid], the first unknown element:

• If it is 0, it belongs at the end of the zeros region. Swap it with nums[low], then advance both low and mid. The element that was at nums[low] was a 1 (or the same element if low == mid), so after the swap nums[mid] is a known 1 and it is safe to move past it.
• If it is 1, it is already in the right place: just advance mid.
• If it is 2, swap it with nums[high] and shrink high. Do not advance mid: the element that came from nums[high] is unknown and must be examined next.

Each iteration shrinks the unknown region by one, so after at most n steps it is empty and the three coloured regions are contiguous and in order.

Algorithm.

1. Set low = 0, mid = 0, high = n - 1.
2. While mid ≤ high:
   • nums[mid] == 0: swap nums[low] and nums[mid]; low++, mid++.
   • nums[mid] == 1: mid++.
   • nums[mid] == 2: swap nums[mid] and nums[high]; high--.

static void optimal(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] == 1) {
            mid++;
        } else {
            int t = nums[mid]; nums[mid] = nums[high]; nums[high] = t;
            high--;
        }
    }
}

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

Complexity. Time O(n): every iteration either advances mid or decreases high, so there are at most n iterations, each O(1). Space O(1): three indices and a temporary.

Why is the loop condition mid <= high and not mid < high? When mid == high there is still exactly one unknown element, and it may be a 0 that needs to move left. On [1, 0] the first iteration only advances mid to 1, which equals high; the iteration that swaps the 0 to the front happens only because of the <=. With < the method would return [1, 0].

Edge cases and pitfalls

• Empty array. high = -1, the loop does not run, and nothing breaks. The counting version writes nothing.
• Single element or all equal values: mid walks to the end (for 1s) or the regions grow until they meet; no swaps go out of bounds because low ≤ mid and high ≥ mid hold throughout.
• Do not advance mid after swapping a 2. The element swapped in from the right is unexamined. Advancing anyway is the single most common bug and leaves 2s in the middle.
• Advancing mid after swapping a 0 is correct. Everything in [low, mid) is a 1, so the element swapped in is a 1 (or mid == low and the element is the same 0). Either way it is already sorted.
• Values outside {0, 1, 2}. The counting version lumps anything that is not 0 or 1 into the 2 bucket and would silently corrupt data; the flag version does the same via its else. Validate input if it is not guaranteed.
• Stability does not matter: equal values are indistinguishable, so swapping freely is fine.
• Loop condition. mid <= high, not <; see the paragraph above.

Related problems

• Sort a Linked List of 0’s 1’s and 2’s is the same problem on a linked list, where the counting approach is natural and the swap approach is replaced by splicing three sublists.
• Majority Element-I is another array problem where exploiting a structural property (here, a value with more than half the votes) replaces sorting.
• Kadane’s Algorithm is the next problem on the sheet and another one-pass, constant-space sweep.