Adrian Wang's blog

一个完全由AI生成，但是内容高质量的Blog

LeetCode 75 - 颜色分类 (Sort Colors)

Posted on Edited on In , ,
Symbols count in article: 4.2k Reading time ≈ 4 mins.
LeetCode 75 颜色分类题解，使用双指针原地排序包含0、1、2的数组。

问题描述

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

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

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

示例 1：

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

示例 2：

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

约束条件：

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

解题思路

想象一下，我们手里有一堆混杂着红色（0）、白色（1）和蓝色（2）的珠子，目标是把它们排成红、白、蓝的顺序，而且只能在原来的位置上调整。

我们可以用三个"标记"来帮助我们：

  1. left 标记：它指向数组最左边，表示"这里应该是红色珠子（0）的地盘"。所有在 left 左边的珠子保证都是红色的。初始时，left 在数组的第一个位置（索引 0）。
  2. right 标记：它指向数组最右边，表示"这里应该是蓝色珠子（2）的地盘"。所有在 right 右边的珠子保证都是蓝色的。初始时，right 在数组的最后一个位置。
  3. current 标记（在代码中我们用变量 i）：这是我们当前正在检查的珠子。我们从左到右一个一个地看这些珠子。初始时，current 和 left 一样，都在数组的第一个位置。

接下来，我们开始移动 current 标记，检查它指向的珠子是什么颜色：

  • 如果 nums[current] 是红色 (0)：

    • 太好了！红色珠子就应该在左边。我们把它和 left 标记指向的珠子交换位置。
    • 这样一来，left 标记左边的区域又多了一个确定的红色珠子，所以 left 向右移动一位。
    • current 标记也向右移动一位，去看下一个珠子。因为我们是从左边换过来的，所以换到 current 位置的珠子不可能是蓝色（2），只可能是红色（0）或白色（1），这些都是 current 左边已经处理过的，所以可以直接继续。

  • 如果 nums[current] 是蓝色 (2)：

    • 蓝色珠子应该在最右边。我们把它和 right 标记指向的珠子交换位置。
    • 这样一来，right 标记右边的区域又多了一个确定的蓝色珠子，所以 right 向左移动一位。
    • 重点来了： current 标记这次先不移动！为什么呢？因为从 right那边换过来的珠子颜色是不确定的（可能是红、白、蓝任何一种），我们需要在下一轮循环中重新检查一下 current 当前位置这个新换来的珠子是什么颜色。为了抵消 for 循环末尾 current (即 i)会自动增加的效果，我们会让 i 减 1。

  • 如果 nums[current] 是白色 (1)：

    • 白色珠子就应该在红色和蓝色之间。它现在的位置挺好，不需要动。
    • current 标记向右移动一位，去看下一个珠子。

我们一直重复这个过程，直到 current 标记超过了 right 标记。当 current > right 时，就说明所有的珠子都已经各就各位了！

这种方法很巧妙，只需要看一遍所有的珠子（时间复杂度 $O(n)$），而且不需要额外的空间来存放珠子（空间复杂度 $O(1)$）。

让我们通过一个例子来理解这个过程：nums = [2,0,2,1,1,0]

初始状态: left = 0, right = 5 (数组最后一个索引), current = 0
nums = [2, 0, 2, 1, 1, 0]
^ ^
curr (i) right
left

  1. current = 0, nums[0] = 2 (蓝色)。

    • 是蓝色！和 nums[right] (即 nums[5], 值为 0) 交换。
    • 数组变为 [0, 0, 2, 1, 1, 2]
    • right 左移到索引 4。
    • current (即 i) 保持在 0 (通过 i-- 实现)，因为换过来的 nums[0] (现在是 0) 需要重新检查。
      nums = [0, 0, 2, 1, 1, 2]
      ^ ^
      curr right
      left

  2. current = 0, nums[0] = 0 (红色)。

    • 是红色！和 nums[left] (即 nums[0], 值为 0) 交换 (原地不动)。
    • 数组仍为 [0, 0, 2, 1, 1, 2]
    • left 右移到索引 1。
    • current (即 i) 右移到索引 1。
      nums = [0, 0, 2, 1, 1, 2]
      ^ ^
      curr right
      left

  3. current = 1, nums[1] = 0 (红色)。

    • 是红色！和 nums[left] (即 nums[1], 值为 0) 交换 (原地不动)。
    • 数组仍为 [0, 0, 2, 1, 1, 2]
    • left 右移到索引 2。
    • current (即 i) 右移到索引 2。
      nums = [0, 0, 2, 1, 1, 2]
      ^ ^
      curr right
      left

  4. current = 2, nums[2] = 2 (蓝色)。

    • 是蓝色！和 nums[right] (即 nums[4], 值为 1) 交换。
    • 数组变为 [0, 0, 1, 1, 2, 2]
    • right 左移到索引 3。
    • current (即 i) 保持在 2 (通过 i-- 实现)，因为换过来的 nums[2] (现在是 1) 需要重新检查。
      nums = [0, 0, 1, 1, 2, 2]
      ^ ^
      curr right
      left

  5. current = 2, nums[2] = 1 (白色)。

    • 是白色！位置正确，不动。
    • current (即 i) 右移到索引 3。
      nums = [0, 0, 1, 1, 2, 2]
      ^ ^
      curr right
      left

  6. current = 3, nums[3] = 1 (白色)。

    • 是白色！位置正确，不动。
    • current (即 i) 右移到索引 4。
      nums = [0, 0, 1, 1, 2, 2]
      ^
      curr,right
      left

  7. current = 4。此时 current (4) > right (3)。循环结束。

最终排序结果: [0,0,1,1,2,2]

实现细节

代码的逻辑和我们上面描述的"珠子排序"思路是完全一样的。

  • left 变量用来标记红色珠子（0）区域的右边界。
  • right 变量用来标记蓝色珠子（2）区域的左边界。
  • 循环变量 i 就是我们的 current 标记，它从数组开头（索引 0）一直扫描到 right 标记的位置。

在 for 循环里面：

  • 如果 nums[i] 是 0（红色）：就把它跟 nums[left] 交换，然后 left 和 i 都向右挪一个位置。
  • 如果 nums[i] 是 2（蓝色）：就把它跟 nums[right] 交换，然后 right 向左挪一个位置。特别注意： 这时候 i 要减 1 (i--)，这是因为从 right那边换过来的珠子我们还没看过是啥颜色，通过 i-- 可以抵消掉 for 循环自己会做的 i++，这样下一轮循环就能重新检查当前 i 位置的这个新珠子了。
  • 如果 nums[i] 是 1（白色）：那它就在正确的位置（红和蓝的中间），我们啥也不用做，i 会在 for 循环结束时自动向右挪。

代码实现

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
30
31
32
33
package main

/*
 * @lc app=leetcode.cn id=75 lang=golang
 *
 * [75] 颜色分类
 */

// @lc code=start
func sortColors(nums []int) {
	left, right := 0, len(nums)-1 // 初始化 left 和 right 标记

	// current 标记 i 从数组头扫描到 right 标记的位置
	for i := 0; i <= right; i++ {
		if nums[i] == 0 { // 遇到红色珠子 (0)
			// 和 left 位置的珠子交换
			nums[i], nums[left] = nums[left], nums[i]
			// left 标记向右移动
			left++
			// current 标记 (i) 也会在 for 循环末尾自动右移
			// 因为从 left 换过来的珠子只可能是 0 或 1, 已经处理过或即将正确处理
		} else if nums[i] == 2 { // 遇到蓝色珠子 (2)
			// 和 right 位置的珠子交换
			nums[i], nums[right] = nums[right], nums[i]
			// right 标记向左移动
			right--
			// 关键：current 标记 (i) 退一步，重新检查当前位置换过来的新珠子
			i--
		}
		// 如果 nums[i] == 1 (白色珠子)，啥也不做，i 会自动右移
	}
}
// @lc code=end

方法比较

方面 双指针法 (荷兰国旗问题解法) 计数排序法
时间复杂度 $O(n)$ $O(n)$ (两次遍历)
空间复杂度 $O(1)$ (原地) $O(k)$ (k=3, 颜色数量)
优点 原地排序，空间效率高 实现简单直观
缺点 逻辑稍复杂 需要额外空间
推荐度 ★★★★★ ★★★★☆

计数排序法思路：

  1. 遍历数组，统计 0, 1, 2 各自出现的次数。
  2. 根据统计的次数，重写原数组。先写所有的 0，然后写所有的 1，最后写所有的 2。

这种方法也很简单，但题目要求原地排序，并且如果颜色种类很多，计数排序的空间复杂度会上升。对于本题，颜色只有三种，所以空间复杂度是 $O(1)$。

复杂度分析

双指针法：

  • 时间复杂度： $O(n)$。
    left 和 current 指针从左向右移动，right 指针从右向左移动。每个元素最多被访问和交换常数次。current 指针最多遍历数组一次。
  • 空间复杂度： $O(1)$。
    我们只使用了常数个额外的变量（left, right, current），是在原地修改数组。

关键收获

  • 双指针（或三标记）是个好技巧：对于数组排序或分区问题，这种方法非常高效。这个特定的问题也叫"荷兰国旗问题"，因为荷兰国旗就是红、白、蓝三色条纹。
  • 原地排序省地方：我们没有用额外的数组，直接在原来的数组上操作，这很节省内存。
  • 指针移动要小心：
    • 当遇到蓝色珠子（2）并与右边交换后，current (即 i) 指针需要"退一步" (i--) 是这个算法的关键。这能确保我们不会错过检查从右边换过来的那个未知颜色的珠子。
    • 当遇到红色珠子（0）并与左边交换后，current (即 i) 和 left 都前进。这是因为 left 左边的区域保证都是红色，而 current 是从左往右扫描的，所以从 left 换到 current 位置的珠子不可能是蓝色（2）。它只可能是红色（0）或白色（1）。如果换过来的是红色，current 下次检查时会再次把它放到 left区域；如果换过来的是白色，current 下次检查时会直接跳过。

总的来说，这道题通过巧妙的指针操作，实现了高效的颜色排序，是一个值得学习的经典算法。