Leetcode: Sort Colors

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.

click to show follow up.

Follow up:
A rather straight forward solution is a two-pass algorithm using counting sort.
First, iterate the array counting number of 0's, 1's, and 2's, then overwrite array with total number of 0's, then 1's and followed by 2's.

Could you come up with an one-pass algorithm using only constant space?

Understand the problem:
The problem reads likes an application problem. However, if you can abstract from it, it actually gives you an arrays of integers, with the number 0, 1, and 2. For example, 1, 1, 2, 0, 1, 0, 1. Then you are asked to sort it out. Note that you are not allowed to use any sorting libraries like quick sort. 

Naive Solution:
Since you know that the array is composed by only 0s, 1s and 2s. You should naturally think about the bucket sort. There are three buckets, of each contains 0s, 1s, and 2s. The most straight-forward solution is first to count the number of 0s, 1s and 2s respectively. Then in the second pass fill the original array with the number of 0, 1, and 2. 

Code (Java):

public class Solution {
    public void sortColors(int[] A) {
        if (A == null || A.length <= 1) {
            return;
        }
        
        int num0 = 0;
        int num1 = 0;
        int num2 = 0;
        
        for (int i = 0; i < A.length; i++) {
            switch(A[i]) {
                case 0: ++num0;
                        break;
                case 1: ++num1;
                        break;
                case 2: ++num2;
                        break;
            }
        }
        
        int j = 0;
        for (j = 0; j < num0; j++) {
            A[j] = 0;
        }
        
        for (j = 0; j < num1; j++) {
            A[num0 + j] = 1;
        }
        
        for (j = 0; j < num2; j++) {
            A[num0 + num1 + j] = 2;
        }
    }
}

A better approach:
The question asks for solving the problem using only one pass of the array. Since the array is only composed of 0, 1, and 2, we can naturally think of using two pointers approach. We use two pointers, red and blue, points to the starting and end of the array initially. Then iterate index i from 0 to the end of the array. If A[i] == 0, move it to red pointer. If A[i] == 2, move it to blue pointer. Else move on. 

Code (Java):

public class Solution {
    public void sortColors(int[] A) {
        if (A == null || A.length <= 1) {
            return;
        }
        
        int red = 0;
        int blue = A.length - 1;
        
        int i = 0;
        while (i <= blue) {
            if (A[i] == 0) {
                swap(A, i, red);
                red++;
                i++;
            } else if (A[i] == 2) {
                swap(A, i, blue);
                blue--;
            } else {
                i++;
            }
        }
    }
    
    private void swap(int[] A, int i, int j) {
        int temp = A[i];
        A[i] = A[j];
        A[j] = temp;
    }
}

Discussion:
Note that the increase of index i. For A[i] == 0, we must increase the index of i once we swap A[i] and A[red], because i cannot be less than red. For the A[i] == 2, we don't need to increase i, since now A[i] could be possible 0 since it is swapped with 2. We need to check that number.

Summary:
This problem is actually a 3-way partition problem.