75. 颜色分类

LeetCode 75¹ · 中等 · 技巧

题

数组 nums 只含 0、1、2 三种值（代表红白蓝），原地把它排成所有 0 在前、1 居中、2 在后，不许调用库排序。进阶：一趟扫描、常数空间。

例：nums = [2, 0, 2, 1, 1, 0] → 原地变为 [0, 0, 1, 1, 2, 2]。

一句话

两个边界指针分别收 0 和 2，中间指针扫描：遇 0 扔到左边、遇 2 扔到右边、遇 1 跳过。

关键技巧

用循环不变量写指针题。 扫描过程中始终保持：

• nums[0 : p0] 全是 0；
• nums[p0 : i] 全是 1；
• nums[i : p2 + 1] 未处理；
• nums[p2 + 1 : n] 全是 2。

i > p2 时未处理区为空，排序完成。每一步只需问：处理 nums[i] 后怎样维持这四段？

三种情况。

• nums[i] == 0：和 nums[p0] 交换。换过来的一定是 1（来自已处理的 1 区），或者 p0 == i 时是它自己——所以 p0、i 都前进。
• nums[i] == 1：已在正确位置，i 前进。
• nums[i] == 2：和 nums[p2] 交换，p2 后退；换过来的是未处理区的元素，没看过，所以 i 不动，下一轮再判。

这是快排三路划分的特例。 以 1 为 pivot 的 partition——小于、等于、大于各一段。

朴素做法对比。 计数排序：数出 0、1、2 的个数再覆写，两趟扫描，也是 \(O(1)\) 空间。

解

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        p0, i, p2 = 0, 0, len(nums) - 1
        while i <= p2:
            if nums[i] == 0:
                nums[p0], nums[i] = nums[i], nums[p0]
                p0 += 1
                i += 1
            elif nums[i] == 2:
                nums[p2], nums[i] = nums[i], nums[p2]
                p2 -= 1          # 换来的元素还没看过，i 不动
            else:
                i += 1

时间 \(O(n)\)，一趟；空间 \(O(1)\)。

延伸

• 两路版：283. 移动零，只分「非 0 / 0」两段，一个写指针就够。
• 三路 partition 的用处：215. 数组中的第K个最大元素 的快速选择遇到大量重复元素时，三路划分能避免退化到 \(O(n^2)\)。
• 坑：遇 2 交换后 i += 1 会漏判换来的 0 或 2——[2, 0, 1] 就能测出来。

s = Solution()
for arr, want in [
    ([2, 0, 2, 1, 1, 0], [0, 0, 1, 1, 2, 2]),
    ([2, 0, 1], [0, 1, 2]),
    ([0], [0]),
    ([2, 2, 1, 1, 0, 0], [0, 0, 1, 1, 2, 2]),
]:
    s.sortColors(arr)
    assert arr == want