75.颜色分类

Bin大约 5 分钟

题目描述

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

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

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

进阶： 你能想出一个仅使用常数空间的一趟扫描算法吗？

示例 1:

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

解题思路：双指针 (荷兰国旗问题)

这道题是经典的“荷兰国旗问题”。要求原地排序，并且最好只遍历一次，空间复杂度为 O(1)。

我们可以想象将数组分成三个部分：

[0, p0 - 1]：全部是 0 (红色)
[p0, i - 1]：全部是 1 (白色)
[p2 + 1, n - 1]：全部是 2 (蓝色)

我们使用三个指针来维护这三个区域的边界：

p0：指向第一个 1 的位置，p0 左边（不含 p0）全是 0。
i：当前遍历到的元素位置。
p2：指向最后一个 1 的位置，p2 右边（不含 p2）全是 2。

算法流程：

初始化 p0 = 0, i = 0, p2 = n - 1。
循环遍历数组，当 i <= p2 时：
- 如果 nums[i] == 0：
  - 当前元素是红色，应该放到 p0 的位置。
  - 交换 nums[i] 和 nums[p0]。
  - p0 和 i 都向右移动一位。因为从 p0 换过来的元素我们之前已经检查过了（它要么是 0 要么是 1），所以 i 可以安全地前进。
- 如果 nums[i] == 1：
  - 当前元素是白色，它就在它应该在的区域，不需要动。
  - i 向右移动一位。
- 如果 nums[i] == 2：
  - 当前元素是蓝色，应该放到 p2 的位置。
  - 交换 nums[i] 和 nums[p2]。
  - p2 向左移动一位。注意：此时 i 不能动，因为从 p2 换过来的元素 nums[i] 是一个新元素，我们还没有检查过它，需要下一轮循环来处理它。
当 i 越过 p2 时，排序完成。

动画演示

以 nums = [2, 0, 2, 1, 1, 0] 为例：

步骤	`p0`	`i`	`p2`	`nums`	操作
初始	0	0	5	`[2, 0, 2, 1, 1, 0]`	`nums[0]==2`, swap(0,5)
1	0	0	4	`[0, 0, 2, 1, 1, 2]`	`nums[0]==0`, swap(0,0), p0++, i++
2	1	1	4	`[0, 0, 2, 1, 1, 2]`	`nums[1]==0`, swap(1,1), p0++, i++
3	2	2	4	`[0, 0, 2, 1, 1, 2]`	`nums[2]==2`, swap(2,4)
4	2	2	3	`[0, 0, 1, 1, 2, 2]`	`nums[2]==1`, i++
5	2	3	3	`[0, 0, 1, 1, 2, 2]`	`nums[3]==1`, i++
6	2	4	3	`[0, 0, 1, 1, 2, 2]`	`i > p2`, 循环结束

最终得到 [0, 0, 1, 1, 2, 2]。

代码实现

import java.util.Scanner;

public class Main {
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);

        // 读取数组长度
        int n = scanner.nextInt();
        int[] nums = new int[n];

        // 读取数组元素
        for (int i = 0; i < n; i++) {
            nums[i] = scanner.nextInt();
        }

        Solution solution = new Solution();
        solution.sortColors(nums);

        // 输出排序后的数组
        for (int i = 0; i < n; i++) {
            System.out.print(nums[i]);
            if (i < n - 1) System.out.print(" ");
        }
        System.out.println();

        scanner.close();
    }
}

class Solution {
    public void sortColors(int[] nums) {
        int p0 = 0;          // 指向 0 区域的末尾
        int i = 0;           // 当前遍历指针
        int p2 = nums.length - 1; // 指向 2 区域的开头

        while (i <= p2) {
            if (nums[i] == 0) {
                swap(nums, i, p0);
                p0++;
                i++;
            } else if (nums[i] == 1) {
                i++;
            } else { // nums[i] == 2
                swap(nums, i, p2);
                p2--;
                // 注意：i 不增加，因为换过来的 nums[i] 需要重新判断
            }
        }
    }

    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

复杂度分析

时间复杂度：O(n)
- i 指针从头到尾扫描一次数组。
空间复杂度：O(1)
- 只使用了 p0, i, p2 三个指针，是常量级的额外空间。

解题思路二：单指针覆盖法

这是另一种非常巧妙的单次遍历解法，它不通过交换，而是通过覆盖来完成排序。

我们使用两个指针 p0 和 p1，分别用来记录下一个 0 和 1 应该被放置的位置。

算法流程：

初始化 p0 = 0, p1 = 0。
遍历数组，对于当前元素 nums[i]：
1. 先将 nums[i] 的值保存到临时变量 x 中。
2. 大胆假设：假设当前位置最终应该放 2，所以先用 2 覆盖 nums[i]。
3. 修正假设：
  - 如果原始值 x 小于等于 1（即 x 是 0 或 1），说明假设错误，当前区域至少应该是个 1。于是在 nums[p1] 的位置放一个 1，然后 p1++。
  - 如果原始值 x 等于 0，说明假设更进一步错误，当前区域不仅不是 2，甚至不是 1，而应该是 0。于是在 nums[p0] 的位置放一个 0，然后 p0++。

这个过程的精髓在于，当遇到一个 0 时，它会先在 p1 处放一个 1，然后立即在 p0 处（此时 p0 <= p1）用 0 覆盖掉刚才的 1（或者更早之前放的 1），从而保证了 0 总是在 1 的前面。

代码实现

class Solution {
    public void sortColors(int[] nums) {
        int p0 = 0;
        int p1 = 0;
        for(int i = 0; i < nums.length; i++) {
            int num = nums[i];
            nums[i] = 2;
            if(num <= 1) nums[p1++] = 1;
            if(num == 0) nums[p0++] = 0;
        }
    }
}

复杂度分析

时间复杂度：O(n)，单次遍历。
空间复杂度：O(1)，仅使用常数个额外指针。

全变成2 记录并赋值的1个数记录并赋值的0个数这样就是 0 1 2排序了