75. 颜色分类
题目描述
输入长度为 n 的数组 nums，每个元素是 0、1 或 2，分别代表红、白、蓝三种颜色。
在原数组上重新排列，让同色元素挨在一起，并按红、白、蓝的顺序排列，即所有 0 在前，然后是所有 1，最后是所有 2。
不能调用库里现成的排序函数。直接修改 nums，不需要返回值。
示例 1
输入：nums = [2,0,2,1,1,0] 输出：[0,0,1,1,2,2]
解释：两个 0、两个 1、两个 2 依次排好。
示例 2
输入：nums = [2,0,1] 输出：[0,1,2]
解释：0、1、2 各一个，排成 [0,1,2]。
数据范围
n == nums.length1 <= n <= 300nums[i]只会是0、1或2。
进阶：能否只扫描一遍，并且只用常数额外空间？
解题思路：三指针（荷兰国旗）
荷兰国旗三指针：low 左侧全是 0，mid 是待处理位置，high 右侧全是 2。
为什么正确
[0,low) 为 0，[low,mid) 为 1，(high,n) 为 2，mid 到 high 是未知区。
复杂度分析
时间 O(n)，空间 O(1)。
易错点
遇到 2 与 high 交换后不要立刻增加 mid，因为换回来的元素还没有检查。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。
from typing import List
class Solution:
def sortColors(self, nums: List[int]) -> None:
low = mid = 0
high = 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], nums[high] = nums[high], nums[mid]
high -= 1class Solution {
public void sortColors(int[] nums) {
int low = 0, mid = 0;
int high = nums.length - 1;
while (mid <= high) {
if (nums[mid] == 0) {
int t = nums[low]; nums[low] = nums[mid]; nums[mid] = t;
low++;
mid++;
} else if (nums[mid] == 1) {
mid++;
} else {
int t = nums[mid]; nums[mid] = nums[high]; nums[high] = t;
high--;
}
}
}
}class Solution {
public:
void sortColors(vector<int>& nums) {
int low = 0, mid = 0;
int high = nums.size() - 1;
while (mid <= high) {
if (nums[mid] == 0) {
swap(nums[low], nums[mid]);
low++;
mid++;
} else if (nums[mid] == 1) {
mid++;
} else {
swap(nums[mid], nums[high]);
high--;
}
}
}
};func sortColors(nums []int) {
low, mid := 0, 0
high := len(nums) - 1
for mid <= high {
if nums[mid] == 0 {
nums[low], nums[mid] = nums[mid], nums[low]
low++
mid++
} else if nums[mid] == 1 {
mid++
} else {
nums[mid], nums[high] = nums[high], nums[mid]
high--
}
}
}var sortColors = function(nums) {
let low = 0, mid = 0;
let high = nums.length - 1;
while (mid <= high) {
if (nums[mid] === 0) {
[nums[low], nums[mid]] = [nums[mid], nums[low]];
low++;
mid++;
} else if (nums[mid] === 1) {
mid++;
} else {
[nums[mid], nums[high]] = [nums[high], nums[mid]];
high--;
}
}
};function sortColors(nums: number[]): void {
let low = 0, mid = 0;
let high = nums.length - 1;
while (mid <= high) {
if (nums[mid] === 0) {
[nums[low], nums[mid]] = [nums[mid], nums[low]];
low++;
mid++;
} else if (nums[mid] === 1) {
mid++;
} else {
[nums[mid], nums[high]] = [nums[high], nums[mid]];
high--;
}
}
};其他解法
计数排序两遍扫描、逻辑最简单；三指针一遍扫描并原地交换完成划分，是快速排序三路划分的核心技巧。
解法 2：计数排序 时间 O(n)，空间 O(1)
颜色只有 0、1、2 三种。第一遍统计每种颜色出现的次数，第二遍按 0、1、2 的顺序依次把对应个数写回数组。
代码最简单、不易出错，但需要遍历两遍，不满足进阶要求的“一趟扫描”。只适用于取值种类很少的情况。
from typing import List
class Solution:
def sortColors(self, nums: List[int]) -> None:
cnt = [0, 0, 0]
for x in nums:
cnt[x] += 1
k = 0
for color in range(3):
for _ in range(cnt[color]):
nums[k] = color
k += 1class Solution {
public void sortColors(int[] nums) {
int[] cnt = new int[3];
for (int x : nums) cnt[x]++;
int k = 0;
for (int color = 0; color < 3; color++)
for (int t = 0; t < cnt[color]; t++) nums[k++] = color;
}
}class Solution {
public:
void sortColors(vector<int>& nums) {
int cnt[3] = {0, 0, 0};
for (int x : nums) cnt[x]++;
int k = 0;
for (int color = 0; color < 3; color++)
for (int t = 0; t < cnt[color]; t++) nums[k++] = color;
}
};func sortColors(nums []int) {
cnt := [3]int{}
for _, x := range nums {
cnt[x]++
}
k := 0
for color := 0; color < 3; color++ {
for t := 0; t < cnt[color]; t++ {
nums[k] = color
k++
}
}
}var sortColors = function(nums) {
const cnt = [0, 0, 0];
for (const x of nums) cnt[x]++;
let k = 0;
for (let color = 0; color < 3; color++)
for (let t = 0; t < cnt[color]; t++) nums[k++] = color;
};function sortColors(nums: number[]): void {
const cnt = [0, 0, 0];
for (const x of nums) cnt[x]++;
let k = 0;
for (let color = 0; color < 3; color++)
for (let t = 0; t < cnt[color]; t++) nums[k++] = color;
}