颜色分类

给定一个包含红色、白色和蓝色，共 n 个元素的数组 nums。

请原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。

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

必须在不使用库内置的 sort 函数的情况下解决这个问题。

示例 1：

输入：nums = [2,0,2,1,1,0]
输出：[0,0,1,1,2,2]
解释：
该数组包含两个 0、两个 1 和两个 2。
将它们原地排序后，所有 0 排在最前面，接着是所有 1，最后是所有 2。

示例 2：

输入：nums = [2,0,1]
输出：[0,1,2]
解释：
数组中有且仅有一个 0、一个 1 和一个 2，按 0、1、2 的顺序原地排列。

提示：

• n == nums.length
• 1 <= n <= 300
• nums[i] 为 0、1 或 2

进阶：

你能想出一个仅使用常数空间的一趟扫描算法吗？

解法一（计数法）：

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int count0 = 0;
        int count1 = 0;
        int count2 = 0;

        for (int num : nums)
        {
            if (num == 0)
            {
                count0++;
            }
            else if (num == 1)
            {
                count1++;
            }
            else
            {
                count2++;
            }
        }

        int index = 0;

        while (count0--)
        {
            nums[index++] = 0;
        }

        while (count1--)
        {
            nums[index++] = 1;
        }

        while (count2--)
        {
            nums[index++] = 2;
        }
    }
};

解法二（三指针一趟扫描）：

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int left = 0;
        int cur = 0;
        int right = nums.size() - 1;

        while (cur <= right)
        {
            if (nums[cur] == 0)
            {
                swap(nums[left], nums[cur]);
                left++;
                cur++;
            }
            else if (nums[cur] == 2)
            {
                swap(nums[cur], nums[right]);
                right--;
            }
            else
            {
                cur++;
            }
        }
    }
};

核心思想

这题的数组中只会出现三种数字：

• 0
• 1
• 2

目标是把所有 0 放到前面，所有 2 放到后面，剩下的 1 自然就在中间。

最直接的做法是先统计 0、1、2 的数量，然后再把数组改写成对应数量的 0、1、2。

这个方法不使用 sort，也只用常数额外空间，但需要两次遍历。

进阶要求一趟扫描，可以使用三指针，也叫荷兰国旗问题。

这题最关键的观察是：

遍历过程中维护三个区域：前面全是 0，后面全是 2，中间还没处理的区域继续扫描。

解法一：计数法

因为数组中只有 0、1、2 三种值，所以可以先数它们分别出现了多少次。

第一次遍历统计：

count0
count1
count2

然后从数组开头开始改写：

• 先写 count0 个 0
• 再写 count1 个 1
• 最后写 count2 个 2

这种方法逻辑清晰。

不过它不是一趟扫描，因为需要先统计，再回填。

解法二：三指针一趟扫描

使用三个指针：

• left：下一个 0 应该放的位置
• cur：当前正在检查的位置
• right：下一个 2 应该放的位置

在扫描过程中，数组被分成四个区域：

[0, left - 1]       全是 0
[left, cur - 1]     全是 1
[cur, right]        还没处理
[right + 1, n - 1]  全是 2

每次只看 nums[cur]。

1. 如果 nums[cur] == 0

0 应该放到左边。

所以交换：

swap(nums[left], nums[cur]);

交换后，left 位置已经放好了一个 0。

同时换到 cur 的元素一定已经处理过，或者就是当前元素自己。

所以：

left++;
cur++;

2. 如果 nums[cur] == 1

1 本来就应该放在中间。

所以不需要交换，直接继续扫描：

cur++;

3. 如果 nums[cur] == 2

2 应该放到右边。

所以交换：

swap(nums[cur], nums[right]);

交换后，right 位置已经放好了一个 2。

然后：

right--;

但此时 cur 不能增加。

因为从 right 换过来的元素还没有被检查过，它可能是 0、1 或 2。

必须继续检查当前位置。

为什么遇到 2 时不能移动 cur

例如：

nums = [1,2,0]

当 cur 指向下标 1 时，nums[cur] == 2。

把它和右边交换后得到：

[1,0,2]

此时换到 cur 位置的是 0。

如果立刻 cur++，这个 0 就不会被放到左边，结果会出错。

所以遇到 2 时，只移动 right，不移动 cur。

等下一轮继续处理换过来的元素。

为什么遇到 0 时可以移动 cur

当 nums[cur] == 0 时，会和 nums[left] 交换。

根据区域定义，[left, cur - 1] 之间全是 1。

所以如果 left < cur，换到 cur 的一定是 1。

如果 left == cur，只是原地交换。

因此交换后，cur 位置已经是处理好的中间区域，可以安全向右移动。

所以遇到 0 时执行：

left++;
cur++;

正确性证明

我们证明：三指针算法结束后，数组按照 0、1、2 的顺序排列。

结论 1：循环过程中，四个区域的含义始终成立

初始时：

left = 0
cur = 0
right = n - 1

此时：

• [0, left - 1] 为空
• [left, cur - 1] 为空
• [cur, right] 是整个数组
• [right + 1, n - 1] 为空

区域含义成立。

每次处理 nums[cur]：

• 如果是 0，就放到左侧 0 区域，并扩大 0 区域
• 如果是 1，就把它纳入中间 1 区域
• 如果是 2，就放到右侧 2 区域，并扩大 2 区域

每种操作都会保持四个区域的定义不变。

结论 2：算法不会漏处理任何元素

循环条件是：

cur <= right

这正好表示未处理区域 [cur, right] 非空。

每次循环都会让未处理区域缩小：

• 遇到 0 或 1 时，cur++
• 遇到 2 时，right--

所以最终未处理区域会变为空。

而遇到 2 时不移动 cur，保证了从右边换过来的元素不会被跳过。

因此算法不会漏处理任何元素。

结论 3：循环结束时，数组已经完成排序

循环结束时：

cur > right

说明未处理区域 [cur, right] 已经为空。

根据结论 1，此时数组只剩三个区域：

• 左侧全是 0
• 中间全是 1
• 右侧全是 2

这正好就是题目要求的顺序。

得出结论

由结论 1 可知，算法维护的区域定义始终正确。

由结论 2 可知，所有元素都会被处理。

由结论 3 可知，处理结束后数组满足 0、1、2 的顺序。

因此三指针算法正确。

举例理解

以：

nums = [2,0,2,1,1,0]

为例。

初始：

left = 0, cur = 0, right = 5

处理过程：

最终得到：

[0,0,1,1,2,2]

复杂度分析

解法一

计数一次，回填一次。

• 时间复杂度：O(n)
• 空间复杂度：O(1)

解法二

每个元素最多被 cur 扫描一次，交换只使用常数额外空间。

• 时间复杂度：O(n)
• 空间复杂度：O(1)

解法二满足进阶要求，是更推荐的做法。