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?

分析：在这里介绍两种方法，一种是two pointers法，另一种基于快速排序中的partition算法。

1.　Two pointers

0应该放在数组的左端，2应该放在数组的右端，我们用两个指针red和blue，分别表示将下一个0或者1移动到的位置，当我们遍历数组时，如果是0就将其放在red位置，如果是2就将其放在blue位置。时间复杂度为O(n), one pass。代码如下：

class Solution {
public:
    void sortColors(int A[], int n) {
        int red = 0, blue = n - 1;
        
        for(int i = 0; i < blue +1;){//i < blue + 1 because the elements after blue are all 2
            if(A[i] == 0)
                swap(A[i++], A[red++]);//only 1 can be in A[red]
            else if(A[i] == 2)
                swap(A[i], A[blue--]);//0 or 1 can be in A[blue]
            else i++;
        }
    }
};

2. 利用快速排序中的partition算法。先将0作为pivot将数组分成两部分，再将1作为pivot将数组partition为两部分，毕。此方法适用于有k种颜色的情况。时间复杂度仍为O(n)，但是two pass。代码如下：

class Solution {
public:
    void sortColors(int A[], int n) {
       partition(partition(A, A + n, bind1st(equal_to<int>(), 0)), A + n,
        bind1st(equal_to<int>(), 1));
    }
};