题目描述

给定一个包含红色、白色和蓝色，一共 n 个元素的数组，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。

此题中，我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。

注意:
不能使用代码库中的排序函数来解决这道题。

样例

示例:
输入: [2,0,2,1,1,0]
输出: [0,0,1,1,2,2]

算法1

(快速排序) $O(nlogk)$

快速排序的变种题，针对这道题，有更优的解法，但是快排做法比较容易扩展至任意多种颜色。

时间复杂度

每次通过颜色将数组分为两份，所以复杂度是$O(nlogk)$

空间复杂度

递归$logk$层，消耗了$O(logk)$的空间

参考文献

C++ 代码

class Solution {
public:
    void sortColors(vector<int>& nums) {
        partition(nums, 0, nums.size() - 1, 0, 2);
    }

    void partition(vector<int> &nums, int start, int end, int lower, int upper){
        if (start >= end || lower >= upper) return;

        int l = start, r = end, pivot = lower + (upper - lower) / 2;
        while (l <= r){
            while (l <= r && nums[l] <= pivot) ++l;
            while (l <= r && nums[r] > pivot) --r;
            if (l <= r) swap(nums[l++], nums[r--]);
        }

        partition(nums, start, r, lower, pivot);
        partition(nums, l, end, pivot + 1, upper);
    }
};

算法2

(计数排序) $O(n)$

就数一下三种颜色各有几个，直接往nums里写就行；
这个做法也很容易推广到多个颜色。

时间复杂度

遍历了2次整个数组，时间复杂度是$O(n)$

空间复杂度

需要$k$个元素记录颜色的数量，所以是$O(k)$

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int n = nums.size(), zero = 0, one = 0, two = 0;
        for (int num: nums){
            if (num == 0) ++zero;
            else if (num == 1) ++one;
            else if (num == 2) ++two;
        }

        for (int i = 0; i < n; ++i){
            if (i < zero) nums[i] = 0;
            else if (i - zero < one) nums[i] = 1;
            else nums[i] = 2;
        }
    }
};

算法3

(双指针) $O(n)$

因为题目规定只有3种颜色，所以我们只需要完成一次类似快排partition的操作即可。

时间复杂度

双指针相遇时遍历了整个数组，时间复杂度是$O(n)$

空间复杂度

双指针使用了$O(1)$的空间

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int l = 0, n = nums.size(), mid = 0, r = n - 1;
        while (mid <= r){
            if (nums[mid] == 0) swap(nums[l++], nums[mid++]);
            else if (nums[mid] == 2) swap(nums[mid], nums[r--]);
            else ++mid;
        }
    }
};