75. 颜色分类

题目描述

思路
看到"相同的颜色相邻"、"原地修改数组"、"对数组原地排序"的字样,我第一眼想到的思路就是使用双指针。在 Hot 100 的双指针专题当中,我们已经接触过「移动零」这样的题目了,这道题其实就是「移动零」当中双指针的应用。
我们使用两种思路来解决这道题目,分别是「两次遍历」和「一次遍历」。
首先是「两次遍历」。两次遍历我们要做什么?第一次遍历我们要做的就是把数组当中的0放到前面;第二次遍历我们要做的就是将数组当中的1放到0的后面。有了这样的思路,我们可以使用快慢指针来实现需求。具体来说,令i是慢指针,j是快指针,如果nums[j] == 0,我们就交换nums[i]和nums[j],并令i += 1。第一次遍历之后,[0, i)的左闭右开区间内都是0。我们接着从i的位置开始找[i + 1, n)子数组当中的1。思路和找0相同,如果遇到nums[j] == 1,就和nums[i]交换,并令i += 1。这样,通过两次遍历,我们可以得到题目中要求的结果。
其次是「一次遍历」。「两次遍历」的思路非常直观,但我们自然能够想到,如果能够建立三个指针p0 / p1 / j,是不是通过一次遍历就可以完成数组的排序?答案是可以的,这也是「题目描述」当中"进阶"部分要求我们实现的思路。具体来说,如果nums[j] == 1,我们直接交换nums[p1]和nums[j]即可,并令p1 += 1。而如果nums[j] == 0,我们不能仅仅简单地交换nums[p0]和nums[j],原因是p0所指的位置可能已经填过1了,因此我们需要判断一道p0和p1的大小关系。如果p0 < p1,说明这个位置之前已经放过1了,由于nums[p0]和nums[j]交换使j的位置放置了1,我们只需要再对nums[p1]和nums[j]进行一次交换即可。最后同时移动p0和p1。这样通过一次遍历,我们就能解决问题。
Golang 题解
「两次遍历」
go
func sortColors(nums []int) {
n := len(nums)
// 双指针, 两次遍历
i := 0
for j := 0; j < n; j ++ {
if nums[j] == 0 {
nums[i], nums[j] = nums[j], nums[i]
i ++
}
}
for j := i; j < n; j ++ {
if nums[j] == 1 {
nums[i], nums[j] = nums[j], nums[i]
i ++
}
}
}
「一次遍历」
go
func sortColors(nums []int) {
n := len(nums)
p0, p1 := 0, 0
for j := 0; j < n; j ++ {
if nums[j] == 1 {
nums[j], nums[p1] = nums[p1], nums[j]
p1 ++
} else if nums[j] == 0 {
nums[j], nums[p0] = nums[p0], nums[j]
if p0 < p1 {
// 如果 p0 < p1, 说明此时 nums[j] 交换得到的数是 1
nums[j], nums[p1] = nums[p1], nums[j]
}
p0 ++
p1 ++
}
}
}