开启掘金成长之旅！这是我参与「掘金日新计划 · 2 月更文挑战」的第14天，点击查看活动详情。
颜色分类
描述
给定一个包含红色、白色和蓝色，一共 n 个元素的数组，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。此题中，我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。
注意: 不能使用代码库中的排序函数来解决这道题。
进阶:
-
一个直观的解决方案是使用计数排序的两趟扫描算法。
首先，迭代计算出 0、1 和 2 元素的个数，然后按照 0、1、2 的排序，重写当前数组。
-
你能想出一个仅使用常数空间的一趟扫描算法吗？
难度
中等
示例
输入: [2, 0, 2, 1, 1, 0]
输出: [0, 0, 1, 1, 2, 2]
思路
对数组进行遍历，使用两个指针分别用来交换 0 和 1。用指针 p0 来交换 0，p1 来交换 1，初始值都为 0。遍历数组时：
当元素值为 1 时，将 nums[i] 和 nums[p1] 交换，然后 p1++。
当元素值为 0 时，将 nums[i] 和 nums[p0] 交换，然后 p0++。p0++ 后，p0 指向的元素可能是 1，后续发生交换会把 1 交换出去，所以当 p0 < p1 时，将 nums[i] 和 nums[p1] 交换，最后无论是否有 p0 < p1，将 p0 和 p1 都后移一个位置。
代码
Rust
pub struct Solution {}
impl Solution {
pub fn sort_colors(nums: &mut Vec<i32>) {
// p0 表示 0 要放置的下标, p1 表示 1 要放置的下标
let (mut p0, mut p1) = (0, 0);
for i in 0..nums.len() {
if nums[i] == 0 {
// 交换 i 和 p0 指向的元素
Solution::swap(nums, i, p0);
// 有可能 p0 指向的元素为 1, 跟 i 发生了交换
if p0 < p1 {
// 将 p1 和 i 交换
Solution::swap(nums, i, p1);
}
p0 += 1;
p1 += 1;
} else if nums[i] == 1 {
// 交换 i 和 p1 指向的元素
Solution::swap(nums, i, p1);
p1 += 1;
}
}
}
fn swap(nums: &mut Vec<i32>, a: usize, b: usize) {
let temp = nums[a];
nums[a] = nums[b];
nums[b] = temp;
}
}
#[test]
fn test_sort_colors() {
let mut nums = vec![2, 0, 2, 1, 1, 0];
println!("{:?}", nums);
Solution::sort_colors(&mut nums);
println!("{:?}", nums);
}
运行结果：
[2, 0, 2, 1, 1, 0]
[0, 0, 1, 1, 2, 2]
Go
func sortColors(nums []int) {
// p0 表示 0 要放置的下标, p1 表示 1 要放置的下标
p0, p1 := 0, 0
for i, num := range nums {
if num == 0 {
// 交换 i 和 p0 指向的元素
nums[i], nums[p0] = nums[p0], nums[i]
// 有可能 p0 指向的元素为 1, 跟 i 发生了交换
if p0 < p1 {
// 将 p1 和 i 交换
nums[i], nums[p1] = nums[p1], nums[i]
}
p0++
p1++
} else if num == 1 {
// 交换 i 和 p1 指向的元素
nums[i], nums[p1] = nums[p1], nums[i]
p1++
}
}
}
func TestSortColors(t *testing.T) {
nums := []int{2, 0, 2, 1, 1, 0}
t.Log(nums)
sortColors(nums)
t.Log(nums)
}
运行结果：
[2 0 2 1 1 0]
[0 0 1 1 2 2]
Java
public class Main {
public static void sortColors(int[] nums) {
// p0 表示 0 要放置的下标, p1 表示 1 要放置的下标
int p0 = 0, p1 = 0;
for (int i = 0; i < nums.length; i++) {
if (nums[i] == 0) {
// 交换 i 和 p0 指向的元素
swap(nums, i, p0);
// 有可能 p0 指向的元素为 1, 跟 i 发生了交换
if (p0 < p1) {
// 将 p1 和 i 交换
swap(nums, i, p1);
}
p0++;
p1++;
} else if (nums[i] == 1) {
// 交换 i 和 p1 指向的元素
swap(nums, i, p1);
p1++;
}
}
}
public static void swap(int[] nums, int a, int b) {
int temp = nums[a];
nums[a] = nums[b];
nums[b] = temp;
}
public static void main(String[] args) {
int[] nums = new int[]{2, 0, 2, 1, 1, 0};
System.out.println(Arrays.toString(nums));
sortColors(nums);
System.out.println(Arrays.toString(nums));
}
}
运行结果：
[2, 0, 2, 1, 1, 0]
[0, 0, 1, 1, 2, 2]
