[LeetCode] 每日一题 75. 颜色分类（双指针）

题目描述

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

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

必须在不使用库内置的 sort 函数的情况下解决这个问题。

题目链接

https://leetcode.cn/problems/sort-colors

示例输入

示例 1

输入：nums = [2,0,2,1,1,0]
输出：[0,0,1,1,2,2]

示例 2

输入：nums = [2,0,1]
输出：[0,1,2]

提示

• n == nums.length

• 1 <= n <= 300

• nums[i] 为 0、1 或 2

解题思路

这题整体思路不算难，很容易联想到用双指针来做，目标是在 O(n) 的时间复杂度下完成排序，且必须是原地修改。

我一开始想到的做法是两次遍历：

• 第一次遍历收集所有的 0，放到前面

• 第二次遍历收集所有的 1，紧接着放在 0 后面
  这种做法虽然时间复杂度也是 O(n)，但效率不够高，空间和实现也不够优雅

后来我想到，其实我们可以用三个指针来优化这个过程，一次遍历就能完成排序：

• redIndex 代表放置 0 的位置（从左向右推进）

• blueIndex 代表放置 2 的位置（从右向左推进）

• 当前遍历指针 i 从头开始遍历整个数组

但是实际写代码时，我发现一个容易踩坑的点是：当遇到 2 并交换时，如果交换回来的是 0，就必须重新处理，否则这个 0 就被跳过了。因此需要在 while 循环里处理 2，而不是简单的 if 判断。

代码实现

class Solution {
    public void sortColors(int[] nums) {
        int n = nums.length;
        int redIndex = 0, blueIndex = n - 1;
        for (int i = 0; i <= blueIndex; i++) {
            while (i <= blueIndex && nums[i] == 2) {
                swap(nums, blueIndex, i);
                blueIndex--;
            }
            if (nums[i] == 0) {
                swap(nums, redIndex, i);
                redIndex++;
            }
        }
    }

    private void swap(int[] nums, int index1, int index2) {
        int temp = nums[index1];
        nums[index1] = nums[index2];
        nums[index2] = temp;
    }
}

复杂度分析

• 时间复杂度：

  • O(n)，只遍历了一次数组，每个元素最多被交换一次

• 空间复杂度：

  • O(1)，使用了常数级别的额外变量，没有额外数组

总结

这道题比较经典，被称为荷兰国旗问题，其实也算是一种分区操作。通过三个指针，我们实现了一次遍历内完成排序，既节省了时间，也保证了空间最优。写的过程中一定要注意交换后的回退处理，是容易出错的地方。

希望这篇分享能为你带来启发！如果你有任何问题或建议，欢迎在评论区留言，与我共同交流探讨。