75. 颜色分类 (Medium)

专题归类: 12-技巧 LeetCode 链接: https://leetcode.cn/problems/sort-colors/


在线做题: 写 Python 代码并运行 · Java 跳转 LeetCode

题目描述

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

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

必须在不使用库的 sort 函数的情况下解决这个问题,且只使用常数级额外空间。

示例 1:

输入:nums = [2,0,2,1,1,0]
输出:[0,0,1,1,2,2]

示例 2:

输入:nums = [2,0,1]
输出:[0,1,2]

题目详细分析

数据范围: n == nums.length1 <= n <= 300nums[i] 取值为 0、1 或 2

核心约束:

  • 原地排序:不能创建新数组
  • 常数空间:O(1) 额外空间
  • 不能使用库函数:不能直接调用 sort()
  • 只有三种值:0、1、2

关键洞察:

  • 因为只有三种值,这比通用的排序要简单得多——本质上是一个三路分区问题
  • 类似快速排序中的 partition 操作,但这里有两个 pivot(0 和 2)
  • 荷兰国旗问题(Dutch National Flag Problem)正是这个问题的原型,由 Dijkstra 提出
  • 可以用三指针法一次遍历完成:用 zero 指针指向 0 区域的右边界,two 指针指向 2 区域的左边界,i 指针遍历数组
  • 也可以先用计数排序:数出 0、1、2 的个数,然后依次覆盖

小白版直白理解

就像整理三种颜色的球:红色、白色、蓝色混在一起,要按红白蓝的顺序排好。

想象三个分区: 左边是红色区、中间是白色区、右边是蓝色区。你拿一个指针从左边走到右边,遇到红球就扔到左边红色区,遇到蓝球就扔到右边蓝色区,遇到白球就留在原地不动。

特别注意的是:当你从右边换回来一个蓝球时,换回来的可能是红球也可能是白球,所以不能急着往前走,要再看一眼换回来的是什么。这就是三指针法的精髓——换回来的球还要再处理一次。


解题思路

思路一:三指针 / 荷兰国旗(推荐)

思路讲解: 设置三个指针:

  • zero:指向 0 区域的右边界(初始为 0)
  • two:指向 2 区域的左边界(初始为 n-1
  • i:当前遍历指针

遍历规则:

  • nums[i] == 0:与 nums[zero] 交换,zero++i++(因为换回来的一定是 1,可以前进)
  • nums[i] == 2:与 nums[two] 交换,two--i 不动,因为换回来的可能是 0 或 1,需要重新检查)
  • nums[i] == 1i++(白色的留在中间,直接跳过)
def sortColors(nums):
    zero = 0          # 0 区域的右边界(下一个 0 应该放的位置)
    i = 0             # 当前遍历指针
    two = len(nums) - 1  # 2 区域的左边界(下一个 2 应该放的位置)
    
    while i <= two:
        if nums[i] == 0:
            # 遇到 0:交换到左边 zero 位置
            nums[i], nums[zero] = nums[zero], nums[i]
            zero += 1
            i += 1
        elif nums[i] == 2:
            # 遇到 2:交换到右边 two 位置
            nums[i], nums[two] = nums[two], nums[i]
            two -= 1
            # 注意:i 不动!因为换回来的可能是 0,需要下一轮继续判断
        else:
            # 遇到 1:留在中间,直接跳过
            i += 1

时间复杂度: O(n) | 空间复杂度: O(1) | 遍历次数: 1 次

思路二:计数排序(两趟扫描)

思路讲解: 先遍历一遍统计 0、1、2 的个数,然后按顺序覆盖原数组。直观易懂,但需要遍历两次。

def sortColors(nums):
    # 第一遍:统计 0、1、2 的个数
    count = [0, 0, 0]
    for num in nums:
        count[num] += 1
    
    # 第二遍:按顺序覆盖
    i = 0
    for color in range(3):
        for _ in range(count[color]):
            nums[i] = color
            i += 1

时间复杂度: O(n) | 空间复杂度: O(1) | 遍历次数: 2 次


易错点

  • 交换 2 时 i 不能前进:这是最常见的 bug。当 nums[i] == 2 并与 nums[two] 交换后,从右边换回来的元素可能是 0 也可能是 1,必须重新判断,所以 i 不能自增
  • 交换 0 时 i 可以前进:因为 nums[zero] 是 0 区域的右边界,nums[zero] 的位置一定是 0 或 1(不会是 2,因为 2 已经被交换到右边了),所以交换后 nums[i] 是 0 或 1,如果是 1 就跳过,如果是 0 会在下一次循环被处理——等一下,实际上当 nums[i]==0 时交换后 nums[zero] 一定是1(因为 zero <= i,且所有 2 已经被交换到 two 之后),所以 i 可以安全地前进
  • 循环条件 while i <= two:当 i > two 时,说明所有位置都已处理好,循环结束。注意是 <= 不是 <
  • 数组长度为 1:三指针法能正确处理,不会进入交换逻辑
  • nums = [2, 0, 1] 测试:这是个好的边界测试用例,可以验证交换逻辑

框架提炼

三指针分区模板:荷兰国旗问题

def threeWayPartition(nums):
    left = 0          # 左分区右边界
    right = len(nums) - 1  # 右分区左边界
    i = 0
    
    while i <= right:
        if nums[i] == LEFT_VALUE:
            nums[i], nums[left] = nums[left], nums[i]
            left += 1
            i += 1
        elif nums[i] == RIGHT_VALUE:
            nums[i], nums[right] = nums[right], nums[i]
            right -= 1
            # i 不前进
        else:  # MIDDLE_VALUE
            i += 1

适用场景: 这个模板适用于将数组按三个类别原地分区的问题,是快速排序三路切分的标准实现。相比于普通的快速排序分区(两路),三路切分在处理大量重复元素时性能更好。

核心要点:

  1. 左指针 left 维护左分区的右边界
  2. 右指针 right 维护右分区的左边界
  3. 遍历指针 i 处理中间区域
  4. 交换右分区时 i 不前进(关键!)

关联题目