排序题目：颜色分类

题目

标题和出处

标题：颜色分类

出处：75. 颜色分类

难度

4 级

题目描述

要求

给定一个包含 $\text{n}$ 个元素的数组 $\text{nums}$，每个元素是红色、白色或蓝色，原地将元素排序，使得相同颜色的元素相邻，并按照红色、白色和蓝色的顺序排列。

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

要求不使用库的排序函数。

示例

示例 1：

输入：$\text{nums = [2,0,2,1,1,0]}$
输出：$\text{[0,0,1,1,2,2]}$

示例 2：

输入：$\text{nums = [2,0,1]}$
输出：$\text{[0,1,2]}$

数据范围

• $\text{n}=\text{nums.length}$
• $\text{1}\le \text{n}\le \text{300}$
• $\text{nums[i]}$ 为 $\text{0}$、$\text{1}$ 或 $\text{2}$

进阶

你可以想出一个只遍历数组一次并使用常数空间的算法吗？

前言

这道题是「荷兰国旗问题」，该问题由荷兰计算机科学家 Edsger Wybe Dijkstra 提出。

解法一

思路和算法

一种直观的解法是使用计数排序的思想，统计数组 $\text{nums}$ 中的元素 $0$、$1$ 和 $2$ 的计数，然后根据计数将元素值按照顺序填入数组 $\text{nums}$，即可完成排序。

代码

class Solution {
    public void sortColors(int[] nums) {
        int[] counts = new int[3];
        for (int num : nums) {
            counts[num]++;
        }
        int n = nums.length;
        for (int i = 0, j = 0; i < n; i++) {
            while (counts[j] == 0) {
                j++;
            }
            nums[i] = j;
            counts[j]--;
        }
    }
}

复杂度分析

• 时间复杂度：$O(n)$，其中 $n$ 是数组 $\text{nums}$ 的长度。需要遍历数组两次，分别统计元素的计数和将元素值按照顺序填入数组。

• 空间复杂度：$O(|\Sigma |)$，其中 $\Sigma$ 是数组 $\text{nums}$ 中的元素集合，这道题中的元素集合只包含 $0$、$1$ 和 $2$，$|\Sigma |=3$。需要记录每个元素的计数。

解法二

思路和算法

解法一需要维护每个元素的计数。也可以在不维护计数的情况下完成排序，需要使用双指针的思想。

由于数组中的元素只有 $0$、$1$ 和 $2$，因此可以遍历数组两次，第一次遍历时将所有的 $0$ 交换到数组的最前面，第二次遍历时将所有的 $1$ 交换到所有的 $0$ 的后面和所有的 $2$ 的前面。

用 $\text{index}$ 表示待填入元素的下标，满足 $\text{index}$ 前面的下标位置都已经填入正确的值。初始时 $\text{index}=0$。

第一次遍历时，如果遇到元素 $0$，则将遇到的元素 $0$ 与下标 $\text{index}$ 处的元素交换，然后将 $\text{index}$ 加 $1$。第一次遍历结束之后，所有的 $0$ 都交换到数组的最前面，$\text{index}$ 指向最后一个 $0$ 的后一个下标位置。

第二次遍历时，如果遇到元素 $1$，则将遇到的元素 $1$ 与下标 $\text{index}$ 处的元素交换，然后将 $\text{index}$ 加 $1$。第二次遍历结束之后，所有的 $1$ 都交换到所有的 $0$ 的后面和所有的 $2$ 的前面。排序结束。

代码

class Solution {
    public void sortColors(int[] nums) {
        int n = nums.length;
        int index = 0;
        for (int i = 0; i < n; i++) {
            if (nums[i] == 0) {
                swap(nums, i, index);
                index++;
            }
        }
        for (int i = index; i < n; i++) {
            if (nums[i] == 1) {
                swap(nums, i, index);
                index++;
            }
        }
    }

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

复杂度分析

• 时间复杂度：$O(n)$，其中 $n$ 是数组 $\text{nums}$ 的长度。需要遍历数组两次，遍历过程中每次交换元素的时间都是 $O(1)$。

• 空间复杂度：$O(1)$。

解法三

思路和算法

解法二需要遍历数组两次。如果使用两个下标分别用于交换两种不同的元素，则可以遍历数组一次完成排序。

使用下标 ${\text{index}}_0$ 和 ${\text{index}}_1$ 分别用于交换元素 $0$ 和元素 $1$，初始时 ${\text{index}}_0$ 和 ${\text{index}}_1$ 都是 $0$。遍历数组的过程中应满足下标 $0$ 到 ${\text{index}}_0-1$ 的元素都是 $0$，下标 ${\text{index}}_0$ 到 ${\text{index}}_1-1$ 的元素都是 $1$。对于遍历到的每个元素，执行如下操作。

• 如果遍历到的元素是 $1$，则将遇到的元素 $1$ 与下标 ${\text{index}}_1$ 处的元素交换，然后将 ${\text{index}}_1$ 加 $1$。

• 如果遍历到的元素是 $0$，则情况较为复杂，需要将遇到的元素 $0$ 与下标 ${\text{index}}_0$ 处的元素交换，交换之后还需要处理已经排序的元素被交换到未排序的部分的情况。用 $i$ 表示遍历到的元素下标，需要执行如下操作。

  1. 将遇到的元素 $0$ 与下标 ${\text{index}}_0$ 处的元素交换。

  2. 如果 ${\text{index}}_0<{\text{index}}_1$，则原来位于 ${\text{index}}_0$ 处的元素是 $1$，该元素 $1$ 被交换到下标 $i$ 处，因此需要将下标 $i$ 处的元素与下标 ${\text{index}}_1$ 处的元素交换。

  3. 将 ${\text{index}}_0$ 和 ${\text{index}}_1$ 分别加 $1$。

遍历结束之后，下标 $0$ 到 ${\text{index}}_0-1$ 的元素都是 $0$，下标 ${\text{index}}_0$ 到 ${\text{index}}_1-1$ 的元素都是 $1$，下标 ${\text{index}}_1$ 到 $n-1$ 的元素都是 $2$，排序结束。

代码

class Solution {
    public void sortColors(int[] nums) {
        int n = nums.length;
        int index0 = 0, index1 = 0;
        for (int i = 0; i < n; i++) {
            if (nums[i] == 1) {
                swap(nums, i, index1);
                index1++;
            } else if (nums[i] == 0) {
                swap(nums, i, index0);
                if (index0 < index1) {
                    swap(nums, i, index1);
                }
                index0++;
                index1++;
            }
        }
    }

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

复杂度分析

• 时间复杂度：$O(n)$，其中 $n$ 是数组 $\text{nums}$ 的长度。需要遍历数组一次，遍历过程中每次交换元素的时间都是 $O(1)$。

• 空间复杂度：$O(1)$。

解法四

思路和算法

解法三在遍历的过程中交换元素 $0$ 和 $1$，也可以在遍历的过程中交换元素 $0$ 和 $2$，同样可以遍历数组一次完成排序。

使用下标 ${\text{index}}_0$ 和 ${\text{index}}_2$ 分别用于交换元素 $0$ 和元素 $2$，初始时 ${\text{index}}_0=0$，${\text{index}}_2=n-1$。遍历数组的过程中应满足下标 $0$ 到 ${\text{index}}_0-1$ 的元素都是 $0$，下标 ${\text{index}}_2+1$ 到 $n-1$ 的元素都是 $2$。对于遍历到的每个元素，执行如下操作。

• 如果遍历到的元素是 $0$，则将遇到的元素 $0$ 与下标 ${\text{index}}_0$ 处的元素交换，然后将 ${\text{index}}_0$ 加 $1$，继续遍历下一个元素。

• 如果遍历到的元素是 $1$，则不执行任何操作，继续遍历下一个元素。

• 如果遍历到的元素是 $2$，则将遇到的元素 $2$ 与下标 ${\text{index}}_2$ 处的元素交换，然后将 ${\text{index}}_2$ 减 $1$。由于在交换元素之后，当前的元素可能变成 $0$，因此不能继续遍历下一个元素，而是应该留在当前元素。

遍历结束之后，下标 $0$ 到 ${\text{index}}_0-1$ 的元素都是 $0$，下标 ${\text{index}}_2+1$ 到 $n-1$ 的元素都是 $2$，所有的元素 $1$ 都位于元素 $0$ 和元素 $2$ 之间，排序结束。

实现方面有一处可以优化。由于遍历过程中确保下标 ${\text{index}}_2$ 右侧的元素都是 $2$，因此当遍历到下标 ${\text{index}}_2$ 右侧时，所有元素都已经排序结束，此时即可结束遍历。

代码

class Solution {
    public void sortColors(int[] nums) {
        int n = nums.length;
        int index0 = 0, index2 = n - 1;
        int i = 0;
        while (i <= index2) {
            if (nums[i] == 0) {
                swap(nums, i, index0);
                index0++;
                i++;
            } else if (nums[i] == 1) {
                i++;
            } else {
                swap(nums, i, index2);
                index2--;
            }
        }
    }

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

复杂度分析

• 时间复杂度：$O(n)$，其中 $n$ 是数组 $\text{nums}$ 的长度。需要遍历数组一次，遍历过程中每次交换元素的时间都是 $O(1)$。

• 空间复杂度：$O(1)$。