LeetCode - 颜色分类

题目链接

leetcode.cn/problems/so…

题目描述

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

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

示例 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

解决方案

先遍历统计0、1、2元素个数，然后给数组赋值；

public void sortColors(int[] nums) {
    if (nums == null || nums.length <= 1) return;
    int zeroNums = 0;
    int oneNums = 0;
    int twoNums = 0;
    for (int n : nums) {
        if (n == 0) {
            zeroNums++;//统计0的个数
        } else if (n == 1) {
            oneNums++;//统计1的个数
        } else if (n == 2) {
            twoNums++;//统计2的个数
        }
    }
    for (int index = 0; index < nums.length; index++) {//给数组赋值
        if (zeroNums > 0) {
            nums[index] = 0 ;
            zeroNums--;
        } else if (oneNums > 0) {
            nums[index] = 1;
            oneNums--;
        } else if (twoNums > 0) {
            nums[index] = 2;
            twoNums--;
        }
    }
}

定义三指针，leftIndex从0开始，交换后保证其左边区间均为元素0，rightIndex从nums.length-1开始，交换后保证其右边区间均为元素2；current从0开始扫描元素；当current指向的元素为2，则与rightIndex交换；指向的元素为0是，则与leftIndex交换；

public void sortColors(int[] nums) {
    if (nums == null || nums.length <= 1) return;
    int leftIndex = 0;
    int currentIndex = 0;
    int rightIndex = nums.length - 1;


    while (currentIndex <= rightIndex) {
        if (nums[currentIndex] == 2) {//当前元素为2，则交换到rightIndex的位置
            int temp = nums[rightIndex];
            nums[currentIndex] = temp;
            nums[rightIndex] = 2;
            rightIndex--;//因为已经保证rightIndex右边为元素2，所以指针往左边走；
            //currentIndex++;当前指针不需要动，因为下一轮要检测当前元素是啥，有可能是0、1、2
        } else if (nums[currentIndex] == 0) {//当前元素为0，则交换到leftIndex的位置
            int temp = nums[leftIndex];
            nums[leftIndex] = 0;
            nums[currentIndex] = temp;
            leftIndex++;
            currentIndex++;//需要往前移动，因为交换后当前位置不是0就是1，不会是2
        } else {//当前元素为1，保留在当前位置
            currentIndex++;
        }
    }
}