【Golang】LeetCode 75. 颜色分类

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 ++
        }
    }
}