排序算法之：利用“快速选择”（Quickselect）算法解决“颜色排序”（Sort Colors）问题

题目描述

给你一个只包含 0， 1 和 2 的整数数组 nums，我们将它们分别视为“红”、“白”、“蓝”三种颜色。请原地对数组进行排序，使得所有 0（红）在前，所有 1（白）在中间，所有 2（蓝）在后。
你不能使用库的排序函数，并且希望在一次遍历中完成，使用常数级的额外空间。

解题过程

这个问题本质上是“荷兰国旗问题”（Dutch National Flag Problem）的一个典型应用。我们可以用“三向划分”的思路来解决，这恰好是“快速选择”（Quickselect）算法核心思想——分区（Partitioning）——的绝佳演练场景。

第一步：理解核心问题与约束

1. 输入限制：数组元素只能是 0, 1, 2。
2. 输出要求：排序后，相同数字必须聚集在一起，顺序为 0, 1, 2。
3. 操作约束：原地排序（空间复杂度 O(1)），并尽量在一次遍历内完成（时间复杂度 O(n)）。

这提示我们不能用基于比较的排序（至少 O(n log n)），也不能用计数排序（需要两次遍历和额外空间）。我们需要一种能“就地”将元素分类到三个区域的方法。

第二步：引入“三向分区”思想

想象数组被分成三个区域，对应三个指针（或索引）来维护边界：

• low 指针：指向最后一个 0 的下一个位置。low 左边的元素（索引小于 low）全部是 0。
• mid 指针：当前正在检查的元素。mid 从左向右移动，将元素分类到正确的区域。
• high 指针：指向第一个 2 的前一个位置。high 右边的元素（索引大于 high）全部是 2。

初始状态：

• 数组未被处理，所以 0 区域和 2 区域都为空。
• 设置 low = 0，mid = 0，high = nums.length - 1。

数组的划分目标就是通过移动这三个指针，最终使得：

• [0, low) 区间：全是 0。
• [low, mid) 区间：全是 1。
• mid 到 high 之间是待处理区域。
• (high, n-1] 区间：全是 2。

当 mid 超过 high 时，排序完成。

第三步：分类规则与指针移动

我们从左到右遍历（用 mid 指针），根据 nums[mid] 的值决定操作：

1. 如果 nums[mid] == 0：

   • 含义：这个 0 应该被放到“0区域”的末尾。
   • 操作：交换 nums[mid] 和 nums[low]。因为 low 指向的是“1区域”的第一个元素（或者就是 mid 自己），交换后，nums[low] 是 0，nums[mid] 变成了原来 low 位置的值（可能是 1 或 0）。
   • 指针移动：
     • low++：因为现在 low 位置已经是 0 了，所以 0 区域的右边界向右扩展一位。
     • mid++：因为交换到 mid 位置的新元素（来自原 low）已经被检查过（它是从 0 区域或 1 区域来的，不可能是 2），所以可以放心地移动 mid 去检查下一个。

2. 如果 nums[mid] == 1：

   • 含义：这个 1 已经在正确的位置（“1区域”）。
   • 操作：什么也不做。
   • 指针移动：mid++。因为 1 应该保留在 low 和 mid 之间，现在 mid 位置是 1，所以只需将 mid 后移，扩大“1区域”。

3. 如果 nums[mid] == 2：

   • 含义：这个 2 应该被放到“2区域”的头部。
   • 操作：交换 nums[mid] 和 nums[high]。因为 high 指向的是“2区域”左边第一个未处理的元素，交换后，nums[high] 变成了 2，而 nums[mid] 变成了原来 high 位置的值（这个值我们尚未检查过）。
   • 指针移动：
     • high--：因为现在 high 位置已经是 2 了，所以 2 区域的左边界向左扩展一位。
     • mid 不动：这是关键！因为交换到 mid 位置的新元素（来自原 high）我们还没有检查过它的值，所以下一轮循环仍需检查当前位置 mid。

第四步：逐步示例

假设 nums = [2, 0, 2, 1, 1, 0]。

初始化：
low = 0, mid = 0, high = 5
数组: [2, 0, 2, 1, 1, 0]
m h
l

步骤分解：

1. nums[mid]=2 -> 交换 mid 和 high, high--。
   数组: [0, 0, 2, 1, 1, 2]
   m h
   l
   （mid 仍为 0，检查新来的 0）

2. nums[mid]=0 -> 交换 mid 和 low, low++, mid++。
   数组: [0, 0, 2, 1, 1, 2]
   m h
   l

3. nums[mid]=0 -> 交换 mid 和 low, low++, mid++。（此时 low 和 mid 都指向索引2）
   数组: [0, 0, 2, 1, 1, 2]
   m h
   l

4. nums[mid]=2 -> 交换 mid 和 high, high--。
   数组: [0, 0, 1, 1, 2, 2]
   m h
   l

5. nums[mid]=1 -> 只做 mid++。
   数组: [0, 0, 1, 1, 2, 2]
   m h
   l

6. nums[mid]=1 -> 只做 mid++。
   此时 mid=5, high=4。循环条件 mid <= high 不再满足，停止。

最终数组: [0, 0, 1, 1, 2, 2]

第五步：算法总结与实现要点

循环不变式（保证算法正确性的关键）：

在每一次循环迭代开始时：

1. [0, low) 区间内的所有元素都是 0。
2. [low, mid) 区间内的所有元素都是 1。
3. (high, n-1] 区间内的所有元素都是 2。
4. mid 指针指向当前待检查的元素。
5. [mid, high] 区间是待处理的、分类未定的元素。

终止条件：当 mid > high 时，所有元素都已处理完毕，并被放到了正确的区域。

代码框架：

def sortColors(nums):
    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
            # 注意：这里不增加 mid

第六步：复杂度分析

• 时间复杂度：O(n)。mid 指针从左到右，high 指针从右到左，每个元素最多被访问一次（交换操作是 O(1)）。
• 空间复杂度：O(1)。只使用了三个整数指针。

这种方法完美满足了题目要求的一次遍历、常数空间、原地排序，是解决此类“有限取值集合排序”问题的经典模板。