Given an array with n objects colored red, white or blue, sort them so that objects of the same color are adjacent, with the colors in the order red, white and blue.

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.

Note:
You are not suppose to use the library’s sort function for this problem.

O(n * 3) in place brute force solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
public class Solution {
    public void sortColors(int[] nums) {
        if (nums != null && nums.length > 1) {
            int left = 0, right = nums.length - 1, colors = 3;
            for (int i = 0; i < colors; i++) {
                while (left <= right) {
                    if (nums[left] != i) {
                        swap(nums, left, right);
                        right--;
                    } else {
                        left++;
                    }
                }
                right = nums.length - 1;
            }
        }
    }
    
    public void swap(int[] nums, int left, int right) {
        int temp = nums[right];
        nums[right] = nums[left];
        nums[left] = temp;
    }
}

Two pass O(n) time O(3) space counting sort solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
public class Solution {
    public void sortColors(int[] nums) {
        if (nums != null && nums.length > 0) {
            int[] count = new int[3];
            for (int num : nums) {
                if (num > 2) 
                    return;
                count[num]++;
            }
            int index = 0;
            for (int i = 0; i < count.length; i++) {
                while (count[i] > 0) {
                    nums[index++] = i;
                    count[i]--;
                }
            }
        }
    }
}

O(n) time O(1) space One pass in place solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
public class Solution {
    public void sortColors(int[] nums) {
        if (nums != null && nums.length > 1) {
            int zero = 0, two = nums.length - 1;
            for (int i = 0; i < nums.length; i++) {
                while (nums[i] == 2 && i < two) {
                    swap(nums, i, two);
                    two--;
                } // must be in this order, test case [1,2,0]
                // if want to check 0 first, the for loop should be in reverse order
                while (nums[i] == 0 && i > zero) {
                    swap(nums, i, zero);
                    zero++;
                }
            }
        }
    }
    
    public void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}