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.

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?

---

We can use two index, one to keep count of red balls and one to keep count of blue balls. And red balls will always be on the left side of the array and blue balls on the right side. We can start by setting red ball index as 0 and blue ball index as n-1. If we find one red ball, swap it with red ball index and keep going. If it is a white ball we don't need to do anything. If it is a blue ball, we can swap it with the blue ball index and keep going. Here is an example:

array: 2, 1, 0 ; red = 0, i = 0, blue = 2;

array: 0, 1, 2 ; red = 0, i = 0, blue = 2;

array: 0, 1, 2 ; red = 0, i = 0, blue = 1;

array: 0, 1, 2 ; red = 1, i = 1, blue = 1;

C++ Code:

/*
 * func: sort_color
 * goal: sort the color array
 * @param A[]: color array
 * @param n: size of the array
 * return:
 */
void sort_color(int A[], int n){
    if(n <= 0){
        return;
    }
    int index = 0;
    int red = 0;
    int blue = n;
    while(index < blue){
        if(A[index] == 0){
            //Put red balls to the front
            swap(A[index++], A[red++]);
        }else if(A[index] == 2){
            //Put blue balls to the end
            swap(A[index], A[--blue]);
        }else{
            ++index;
        }
    }
}

Python Code:

# func: swap two elements in list A
# @param A: input list A
# @param i: index i
# @param j: index j
# @return:
def swap(A, i, j):
    tmp = A[i]
    A[i] = A[j]
    A[j] = tmp

# func: sort the an array only containing 0s, 1s and 2s
# @param A: a list of integers
# @return
# complexity: time O(n), space O(1)
def sort_color(A):
    red = 0
    blue = len(A)-1
    i = 0
    while i <= blue:
        if A[i] == 0:
            swap(A, i, red)
            red += 1
            i += 1
        elif A[i] == 2:
            swap(A, i, blue)
            blue -= 1
        else:
            i += 1