75. 颜色分类#

题目描述#

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

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

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

示例 1：

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

示例 2：

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

提示：

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

解题思路#

核心思想#

数组只含 0/1/2 三种值。我们希望从左到右形成三段：[0 段][1 段][2 段]。用三个指针维护分界：

low：指向下一个 0 该放的位置（0 段的右边界后一位）。
high：指向下一个 2 该放的位置（2 段的左边界前一位）。
mid：当前扫描位置。

0 段在 [0, low)，1 段在 [low, mid)（其实 mid 到 high 之间是未扫描区），2 段在 (high, n-1]。当 mid 扫过 high 时结束。

这是一个经典的 「荷兰国旗问题」，本质是带两个交换锚点的单遍扫描，时间 O(n)、空间 O(1)。

算法步骤#

初始化 low = 0、mid = 0、high = n - 1。
循环直到 mid > high：
- 若 nums[mid] == 0：把它和 nums[low] 交换，然后 low += 1、mid += 1（左边稳定的 0 段扩大）。
- 若 nums[mid] == 2：把它和 nums[high] 交换，然后 high -= 1（右边稳定的 2 段扩大）；不移动 mid，因为换回来的值还没处理。
- 否则（== 1）：mid += 1 直接跳过。

逐步举例#

以 nums = [2, 0, 2, 1, 1, 0]（n=6）为例：

轮次	low	mid	high	数组	动作
初始	0	0	5	[2,0,2,1,1,0]	-
1	0	0	4	[0,0,2,1,1,2]	nums[0]=2 ↔ nums[5]，high→4
2	1	1	4	[0,0,2,1,1,2]	nums[0]=0 ↔ nums[0]，low→1, mid→1
3	2	2	4	[0,0,2,1,1,2]	nums[1]=0 ↔ nums[1]，low→2, mid→2
4	2	2	3	[0,0,1,1,2,2]	nums[2]=2 ↔ nums[4]，high→3
5	2	3	3	[0,0,1,1,2,2]	nums[2]=1，mid→3
6	2	4	3	[0,0,1,1,2,2]	nums[3]=1，mid→4，结束

最终 [0,0,1,1,2,2] ✓

复杂度分析#

时间复杂度：O(n)，mid 每轮至少前进一步。
空间复杂度：O(1)，只用了常数额外空间，原地交换。

代码实现#

1
from typing import List
2

3
class Solution:
4
    def sortColors(self, nums: List[int]) -> None:
5
        """Do not return anything, modify nums in-place instead."""
6
        n = len(nums)
7
        low, mid, high = 0, 0, n - 1
8
        while mid <= high:
9
            if nums[mid] == 0:
10
                nums[low], nums[mid] = nums[mid], nums[low]
11
                low += 1
12
                mid += 1
13
            elif nums[mid] == 2:
14
                nums[mid], nums[high] = nums[high], nums[mid]
15
                high -= 1
16
            else:
17
                mid += 1

1
import sys
2
from pathlib import Path
3
sys.path.insert(0, str(Path(__file__).parent))
4

5
import unittest
6
from solution import Solution
7

8

9
class TestSolution(unittest.TestCase):
10

11
    def test_basic(self):
12
        nums = [2, 0, 2, 1, 1, 0]
13
        Solution().sortColors(nums)
14
        self.assertEqual(nums, [0, 0, 1, 1, 2, 2])
15

16
    def test_three(self):
17
        nums = [2, 0, 1]
18
        Solution().sortColors(nums)
19
        self.assertEqual(nums, [0, 1, 2])
20

21
    def test_already_sorted(self):
22
        nums = [0, 1, 2]
23
        Solution().sortColors(nums)
24
        self.assertEqual(nums, [0, 1, 2])
25

26
    def test_single(self):
27
        nums = [0]
28
        Solution().sortColors(nums)
29
        self.assertEqual(nums, [0])
30

31

32
if __name__ == "__main__":
33
    unittest.main()