LeetCode 75. 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?

分析
由於 0, 1, 2 非常緊湊，首先想到計數排序 (counting sort)，但需要掃描兩遍，不符合題目要求。
由於只有三種顏色，可以設置兩個 index，一個是 red 的 index，一個是 blue 的 index，兩邊往中 間走。時間複雜度 O(n)，空間複雜度 O(1)。
第 3 種思路，利用快速排序裏 partition 的思想，第一次將數組按 0 分割，第二次按 1 分割，排 序完畢，可以推廣到 n 種顏色，每種顏色有重複元素的情況。

代碼

class Solution {
public:
    void sortColors(vector<int>& nums) {
        const int n = nums.size();
        int red = 0;
        int blue = n - 1;

        for (size_t i = 0; i < blue + 1;) {
            if (nums[i] == 0) 
                swap(nums[i++], nums[red++]);
            else if (nums[i] == 2)
                swap(nums[i], nums[blue--]);
            else
                i++;
        }
    }
};