75. 颜色分类

题目链接（中等）

题目描述

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

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

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

数据范围：

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

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

示例

示例 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 的顺序原地排列。

方法一：单指针（两趟扫描）

思路及解法

两次遍历：

1. 第一次遍历把所有 0 交换到数组头部，用指针 ptr 标记头部的右边界；
2. 第二次从 ptr 开始，把所有 1 交换到 0 之后。

遍历结束后，0 在前、1 居中、2 自然都在尾部，排序完成。

代码

class Solution:
    def sortColors(self, nums: list[int]) -> None:
        n = len(nums)
        ptr = 0

        # 第一趟：把 0 换到头部
        for i in range(n):
            if nums[i] == 0:
                nums[i], nums[ptr] = nums[ptr], nums[i]
                ptr += 1

        # 第二趟：把 1 换到 0 之后
        for i in range(ptr, n):
            if nums[i] == 1:
                nums[i], nums[ptr] = nums[ptr], nums[i]
                ptr += 1

复杂度分析

• 时间复杂度：$O(n)$，两趟扫描，每趟都是线性的。
• 空间复杂度：$O(1)$，原地交换，只用常数个变量。

方法二：双指针（p0 和 p1）

思路及解法

用两个指针 p0 和 p1 分别指向「下一个 0 应该放的位置」和「下一个 1 应该放的位置」，均初始为 0。从左到右遍历：

• 遇到 1：与 nums[p1] 交换，p1 += 1；
• 遇到 0：先与 nums[p0] 交换。如果此时 p0 < p1，说明 nums[p0] 原本是 1，被换到了 i 位置，需要再把它与 nums[p1] 交换回去。最后 p0 += 1、p1 += 1。

为什么需要第二次交换？

当 p0 < p1 时，[0, p0) 全是 0，[p0, p1) 全是 1。把 0 换到 p0 时，会把一个 1 挤到 i，如果不处理，这个 1 就留在了后面，答案错误。所以需要再交换一次，把这个 1 放到 p1 位置。

代码

class Solution:
    def sortColors(self, nums: list[int]) -> None:
        n = len(nums)
        p0 = p1 = 0
        for i in range(n):
            if nums[i] == 1:
                nums[i], nums[p1] = nums[p1], nums[i]
                p1 += 1
            elif nums[i] == 0:
                nums[i], nums[p0] = nums[p0], nums[i]
                if p0 < p1:
                    nums[i], nums[p1] = nums[p1], nums[i]
                p0 += 1
                p1 += 1

复杂度分析

• 时间复杂度：$O(n)$，只需一趟扫描。
• 空间复杂度：$O(1)$。

方法三：双指针（p0 和 p2）

思路及解法

用指针 p0 从左边找位置放 0，p2 从右边找位置放 2。遍历指针 i 从左往右走，直到 i > p2 停止。

• 遇到 0：与 nums[p0] 交换，p0 += 1；
• 遇到 2：与 nums[p2] 交换，p2 -= 1。由于换过来的 nums[i] 可能是 0 或 2，需要用 while 反复处理，直到 nums[i] != 2；
• 遇到 1：跳过。

while i <= p2 作为外层循环条件，保证不会越过已排好的 2 区域。

代码

class Solution:
    def sortColors(self, nums: list[int]) -> None:
        n = len(nums)
        p0, p2 = 0, n - 1
        i = 0
        while i <= p2:
            while i <= p2 and nums[i] == 2:
                nums[i], nums[p2] = nums[p2], nums[i]
                p2 -= 1
            if nums[i] == 0:
                nums[i], nums[p0] = nums[p0], nums[i]
                p0 += 1
            i += 1

复杂度分析

• 时间复杂度：$O(n)$，每个元素最多被交换常数次。
• 空间复杂度：$O(1)$。