Sort Colors

NeetCode: Sort Colors · LeetCode: 75. Sort Colors · Difficulty: Medium

Given an array nums with n objects colored red, white or blue (encoded as 0, 1, 2), sort them in place so that objects of the same color are adjacent, in the order red, white, blue. The intended solution does not use the library sort.

Example

Text

Input:  nums = [2, 0, 2, 1, 1, 0]    Output: [0, 0, 1, 1, 2, 2]
Input:  nums = [2, 0, 1]             Output: [0, 1, 2]

Brute force: general-purpose sort

Call the library sort. It is correct but does not exploit that there are only three distinct values. Java arrays have a fixed length, so the method sorts the int[] in place and returns nothing, as LeetCode's signature does.

Java

import java.util.Arrays;

public class Main {
    static class Solution {
        public void sortColors(int[] nums) {
            Arrays.sort(nums);
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        int[] a = {2, 0, 2, 1, 1, 0};
        sol.sortColors(a);
        System.out.println(Arrays.toString(a));
        int[] b = {2, 0, 1};
        sol.sortColors(b);
        System.out.println(Arrays.toString(b));
        int[] c = {};
        sol.sortColors(c);
        System.out.println(Arrays.toString(c));
    }
}

Output

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

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

Better: counting sort, two passes

Count how many zeros, ones and twos there are, then overwrite the array with that many of each. Only three counters are needed.

Java

import java.util.Arrays;

public class Main {
    static class Solution {
        public void sortColors(int[] nums) {
            int[] counts = new int[3];
            for (int x : nums) {
                counts[x]++;
            }
            int i = 0;
            for (int color = 0; color < 3; color++) {
                for (int k = 0; k < counts[color]; k++) {
                    nums[i++] = color;
                }
            }
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        int[] a = {2, 0, 2, 1, 1, 0};
        sol.sortColors(a);
        System.out.println(Arrays.toString(a));
        int[] b = {2, 0, 1};
        sol.sortColors(b);
        System.out.println(Arrays.toString(b));
        int[] c = {};
        sol.sortColors(c);
        System.out.println(Arrays.toString(c));
    }
}

Output

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

Time: O(n) with two passes. Space: O(1).

Optimal: Dutch national flag, one pass

Maintain three regions with pointers low, mid, high: nums[0..low) are all 0, nums[low..mid) are all 1, nums[high+1..] are all 2, and nums[mid..high] (both ends included) is still unknown. Look at nums[mid]:

0: swap with nums[low], then advance both low and mid (the swapped-in value is a known 1).
1: already in place, advance mid.
2: swap with nums[high] and shrink high; do not advance mid because the swapped-in value is unexamined.

high becomes -1 for an empty array or when every element is a 2, which is fine because int can be negative and the loop condition mid <= high then fails.

Java

import java.util.Arrays;

public class Main {
    static class Solution {
        public void sortColors(int[] nums) {
            int low = 0;
            int mid = 0;
            int high = nums.length - 1; // nums[high + 1..] are all 2s
            while (mid <= high) {
                if (nums[mid] == 0) {
                    swap(nums, low, mid);
                    low++;
                    mid++;
                } else if (nums[mid] == 1) {
                    mid++;
                } else {
                    swap(nums, mid, high);
                    high--;
                }
            }
        }

        private void swap(int[] nums, int i, int j) {
            int tmp = nums[i];
            nums[i] = nums[j];
            nums[j] = tmp;
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        int[] a = {2, 0, 2, 1, 1, 0};
        sol.sortColors(a);
        System.out.println(Arrays.toString(a));
        int[] b = {2, 0, 1};
        sol.sortColors(b);
        System.out.println(Arrays.toString(b));
        int[] c = {};
        sol.sortColors(c);
        System.out.println(Arrays.toString(c));
    }
}

Output

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

Time: O(n), each step either advances mid or shrinks the unknown region. Space: O(1).