[LeetCode] 75. Sort Colors（排序颜色）

• Difficulty: Medium

• Related Topics: Array, Two Pointers, Sort

• Link: https://leetcode.com/problems/sort-colors/

Description

Given an array nums with n objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue.
给定一个包含 n 个涂成红、白或蓝色的物体的数组 nums，原地将它们排序，使得相同颜色的物体是相邻的，并且颜色按照红、白、蓝排列。

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.
这里，我们使用整数 0, 1 和 2 分别表示红、白和蓝色。

Follow up

• Could you solve this problem without using the library's sort function?
  你能否不使用语言库自带的排序函数？
• Could you come up with a one-pass algorithm using only O(1) constant space?
  你能想出只使用 O(1) 空间、只遍历一次数组的算法吗？

Examples

Example 1

Input: nums = [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]

Example 2

Input: nums = [2,0,1]
Output: [0,1,2]

Example 3

Input: nums = [0]
Output: [0]

Example 4

Input: nums = [1]
Output: [1]

Constraints

• n == nums.length
• 1 <= n <= 300
• nums[i] is 0, 1, or 2.

Hints

1. A rather straight forward solution is a two-pass algorithm using counting sort.
   一个很直接的想法是两次遍历的计数排序。

2. Iterate the array counting number of 0's, 1's, and 2's.
   遍历数组，统计 0、1 和 2 的数量。

3. Overwrite array with the total number of 0's, then 1's and followed by 2's.
   根据统计到的数量覆写这个数组。

Solution

题目要求不使用语言库自带的排序。当然自己手撸语言库的排序也不是不行，不过提示里给出了一个计数排序的方法，还是相当直观的，并且好像也没有哪个语言内建对计数排序的支持。

class Solution {
    fun sortColors(nums: IntArray): Unit {
        val map = nums.groupBy { it }
            .mapValues { it.value.size }

        var j = 0
        for (i in 0..2) {
            var count = map[i]?:0
            while (count > 0) {
                nums[j++] = i
                count--
            }
        }
    }
}

不过题目里的 Follow up 要求只遍历数组一次，我大概有思路：双指针，遇到 0 往前扔，遇到 2 往后扔，剩下的就是 1 了。不过发现代码不好表达，后面也是从 discussion 里看出来的，双指针不够用就再加一个，代码如下：

class Solution {
    fun sortColors(nums: IntArray): Unit {
        // 三指针，分别表示三种颜色
        var red = 0
        var white = 0
        var blue = nums.lastIndex

        // 因为是按照红、白、蓝进行排序，红和白是前往后走，蓝是后往前走
        // 所以我们主要关注白和蓝的位置
        while (white <= blue) {
            // 主要看 white 指针对应的颜色
            // 红色往前换，蓝色往后换
            when (nums[white]) {
                0 -> {
                    nums.swap(red, white)
                    white++
                    red++
                }
                1 -> {
                    white++
                }
                2 -> {
                    nums.swap(white, blue)
                    blue--
                }
            }
        }
    }

    private fun IntArray.swap(i: Int, j: Int) {
        val t = this[i]
        this[i] = this[j]
        this[j] = t
    }
}