点击获取AI摘要

75. 颜色分类 M

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

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

必须在不使用库内置的 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.length
  • 1 <= n <= 300
  • nums[i] 为 0、1 或 2

问题分析

经典的“Dutch National Flag”问题，需要在一次遍历中将数组分为三部分：

  • 0（红色）
  • 1（白色）
  • 2（蓝色）

算法思路

使用三个指针：

  • low 指向下一个放置 0 的位置；
  • mid 用于扫描当前元素；
  • high 指向下一个放置 2 的位置；

初始时：low = 0，mid = 0，high = n-1。

遍历过程（当 mid <= high）：

  1. 若 nums[mid] == 0，则与 nums[low] 交换，low++, mid++；
  2. 若 nums[mid] == 1，mid++；
  3. 若 nums[mid] == 2，则与 nums[high] 交换，high--（mid 不动，以便新换来的元素再处理）。

这样只需一次扫描，即可将所有元素分类并就地排序。

时间复杂度

  • 时间复杂度：O(n)O(n)，只遍历一次数组
  • 空间复杂度：O(1)O(1)，原地交换，不使用额外空间

代码实现

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

两遍计数 + 原地覆盖

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from typing import List

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        # 1) 计数阶段
        count0 = count1 = count2 = 0
        for v in nums:
            if v == 0:
                count0 += 1
            elif v == 1:
                count1 += 1
            else:
                count2 += 1

        # 2) 原地覆盖阶段
        idx = 0
        # 写入所有 0
        for _ in range(count0):
            nums[idx] = 0
            idx += 1
        # 写入所有 1
        for _ in range(count1):
            nums[idx] = 1
            idx += 1
        # 写入所有 2
        for _ in range(count2):
            nums[idx] = 2
            idx += 1