[leetcode] 75. Sort Colors

题目

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.

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

You must solve this problem without using the library's sort function.

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]

Constraints:

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

Follow up: Could you come up with a one-pass algorithm using only constant extra space?

思路

1. 统计red、white颜色出现的频率，然后填充数组。
2. 遍历序列，遍历的同时保证当前颜色为white，如果不为white，则不断地置换元素直至为white为止。

代码

python版本：

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        cnt1 = nums.count(0)
        cnt2 = nums.count(1)
        nums[:cnt1] = [0]*cnt1
        nums[cnt1:cnt1+cnt2] = [1]*cnt2
        nums[cnt1+cnt2:] = [2]*(len(nums)-cnt1-cnt2)


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