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?

思路

最直白的方法就是利用一趟快速排序的方法，把0，1，2三个数分开。利用到了头尾两个指针。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
public:
    void sortColors(vector<int>& nums) {
        int head = 0;
        while(nums[head] == 0) head++;
        for (int i = head; i < nums.size(); i++){
            int tmp = 0;
            if(nums[i] == 0){
                tmp = nums[head];
                nums[head] = nums[i];
                nums[i] = tmp;
                while(nums[head] == 0) head++;
            }
        }
        int tail = nums.size() - 1;
        while(nums[tail] == 2) tail--;
        for(int i = head; i <= tail; i++){
            int tmp = 0;
            if(nums[i] == 2){
                tmp = nums[tail];
                nums[tail] = nums[i];
                nums[i] = tmp;
                while(nums[tail] == 2) tail--;
            }
        }
    }
};

如果不考虑空间的开销，在题目的提示当中使用的counting sort的方法也是非常好的。数一下0、1、2的个数，然后更新一下原来的数组就可以了。

下面这个方法就厉害了，利用一趟排序、O(1)的空间复杂度就可以完成。We keep a loop invariant that [0,i) [i, j) [j, k) are 0s, 1s and 2s sorted in place for [0,k).

https://discuss.leetcode.com/topic/26181/ac-python-in-place-one-pass-solution-o-n-time-o-1-space-no-swap-no-count