[ 路飞] 颜色分类

题目描述

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

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

解题思路

算法

双指针，p0 表示下一个 0 放置的位置，p2 表示下一个 2 放置的位置

思路

本题的关键在于循环不变量，也就是我们在循环中需要遵循的条件。

首先，在双指针以外，我们需要一个 i 存放遍历的 index

本体的循环不变量是：

[0, p0) = 0
[p0, i] = 1
(p2, r) = 2

其中 r 是我们数组的右边界，在代码中，我们把它替换成 nums.length - 1
可见在遍历中，p0,p2 指向的元素并不是 0，2

我们会在遍历的过程中不断的交换 nums[i], nums[p0] 或者 nums[i], nums[p2]，循环的终止条件是 i <= p2，也就是当我遍历到 2 的区间时，循环应当退出，因此此时 0，1，2 三个值已经都排好了

这个条件的关键点在于是否需要加这个 =

我们来看这么一种情况

2(p0, i), 0, 1(p2)

首先 i 与 p2 的元素交换 ➡️ 1(i, p0), 0(p2), 2

之后 i++ ➡️ 1, 0（i, p2), 2 循环退出

明显可得，此时的数组不和要求，因此在当前我们所设的循环不变量的条件下，循环终止的条件是 i <= p2

代码

/**
 * @param {number[]} nums
 * @return {void} Do not return anything, modify nums in-place instead.
 */
var sortColors = function(nums) {
    let p0 = 0, p2 = nums.length - 1, i = 0

    // loop invariant
    /**
        [0, p0) = 0
        [p0, i] = 1
        (p2, num.length - 1) = 2
     */

    while (i < p2) {
        if (nums[i] === 1) {
            i++
        } else if (nums[i] === 0) {
            ;[nums[p0], nums[i]] = [nums[i], nums[p0]]
            p0++
            i++
        } else if (nums[i] === 2) {
            ;[nums[p2], nums[i]] = [nums[i], nums[p2]]
            p2--
        }
    }
};