算法题--三颜色排序

image.png

0. 链接

题目链接

1. 题目

Given an array 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.

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.

Note: You are not suppose to use the library's sort function for this problem.

Example:

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

Follow up:

A rather straight forward solution is a two-pass algorithm using counting sort.
First, iterate the array counting number of 0's, 1's, and 2's, then overwrite array with total number of 0's, then 1's and followed by 2's.
Could you come up with a one-pass algorithm using only constant space?

2. 思路1：记录三个颜色的位置下标+从左到右遍历

• 初始值设置red, white, blue = 0, 0, len(nums) - 1

• 过程如下图所示

  
  image.png

即跟着white的视角来观察过程

• 当nums[white] == 2时, 需要做一个交换, nums[white], nums[blue] = nums[blue], nums[white], 作用是把当前white所看到的2，转移到数组的BLUE区域的最右边，交换后 blue -= 1, red和white均不变, 因为此时并不清楚换到左边的是什么值, 下一步仍然要检查这个换来的值
• 当nums[white] == 1时, 不需要做修改, 只需要white += 1右移一位, 检查下一个位置, 但此时red并没有变, 也就是说, 这一刻white开始离red而向右而去了, 此时的red可以理解为: 当前RED区域后的第一个位置, 其实是当前WHITE区域的最左位置, 这个位置有什么用呢？它是为了后面可能出现的0准备的
• 当nums[white] == 0时, 注意此时white可能跟red不相同, 那么这个新发现的0, 需要追加到左边RED区域的末尾, 即需要做一个交换, nums[white], nums[red] = nums[red], nums[white], 然后red += 1, white += 1, 表示RED区域扩展了一位, 同时待检查元素也挪到了white之后一位
• 当white == blue时, 说明WHITE区域和BLUE区域会师了，white之后的BLUE区域, 都是2了，不用继续检查了，这样就得到了有序的颜色数组

3. 代码

# coding:utf8
from typing import List


class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        red, white, blue = 0, 0, len(nums) - 1
        while white <= blue:
            if nums[white] == 0:
                nums[red], nums[white] = nums[white], nums[red]
                red += 1
                white += 1
            elif nums[white] == 1:
                white += 1
            else:
                nums[white], nums[blue] = nums[blue], nums[white]
                blue -= 1


def my_test(solution, nums):
    print('input:{}'.format(nums))
    solution.sortColors(nums)
    print('output:{}'.format(nums))
    print('=' * 50)


solution = Solution()
my_test(solution, [2, 0, 2, 1, 1, 0])
my_test(solution, [0, 0, 1, 2, 1, 2])

输出结果

input:[2, 0, 2, 1, 1, 0]
output:[0, 0, 1, 1, 2, 2]
==================================================
input:[0, 0, 1, 2, 1, 2]
output:[0, 0, 1, 1, 2, 2]
==================================================

4. 结果

image.png