题目描述

给定一个包含红色、白色和蓝色、共 n 个元素的数组 nums ，原地 对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。
我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。
必须在不使用库内置的 sort 函数的情况下解决这个问题。

进阶：
你能想出一个仅使用常数空间的一趟扫描算法吗？

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

题目分析

如果允许扫描 2 遍，可以记录每种的个数，然后赋值。
如果只允许扫描 1 遍，可以先思考对于仅 2 种元素的问题如何解决？用1个指针表示第一种元素的预期位置，遍历时遇到第一种元素就和指针进行交换。第二种元素无需特殊处理，自然也会被移动到预期位置。
对于 3 种元素，可以使用 2 个指针，分别表示红色的预期位置（左侧），蓝色的预期位置（右侧），遍历数组时，遇到红色就和红色指针交换，遇到蓝色就和蓝色指针交换。白色无需处理。
需要兼容单个位置交换多次的情况，例如：扫描到了蓝色，与右侧指针交换后恰好是红色。

Java

public void sortColors(int[] nums) {
    int red = 0, blue = nums.length - 1;
    for (int i = 0; i < nums.length; i++) {
        while (i >= red && i <= blue) {
            if (nums[i] == 0) {
                swap(nums, i, red++);
            } else if (nums[i] == 2) {
                swap(nums, i, blue--);
            } else {
                break;
            }
        }
    }
}

private void swap(int[] nums, int i, int j) {
    int temp = nums[i];
    nums[i] = nums[j];
    nums[j] = temp;
}

Kotlin

fun sortColors(nums: IntArray): Unit {
    var red = 0
    var blue = nums.size - 1

    val swap = { i: Int, j: Int ->
        val temp = nums[i]
        nums[i] = nums[j]
        nums[j] = temp
    }

    for (i in nums.indices) {
        while (i in red..blue) {
            when (nums[i]) {
                0 -> swap(i, red++)
                1 -> break
                2 -> swap(i, blue--)
            }
        }
    }
}