Sort Colors II （有k种colors)

题目描述

Given an array of n objects with k different colors (numbered from 1 to k), sort them so that objects of the same color are adjacent, with the colors in the order 1, 2, ... k.

解题方法

brute force

• 先遍历一次记录下每种颜色出现的次数
• 再遍历一次将对应位置改为对应的颜色

这一题就是将sort colors的方法扩展到k种colors，我们首先应该想如何利用本来的sort color的算法

本来sort color是将3种颜色分开，其实是将两种颜色，一种放到开头，一种放到结尾。

我们如果要利用这种方法，就可以每次取出要放在开头的颜色和放在结尾的颜色，也就是Min和max。

将他们按照sort color的方法调整好后，再继续进行下面2个颜色，直到完成所有的颜色

Solution

class Solution:
    """
    @param colors: A list of integer
    @param k: An integer
    @return: nothing
    """
    def sortColors2(self, colors, k):
        # write your code here
        count = 0
        start = 0
        end = len(colors) - 1
        remain = list(colors)
        while count < k:
            if not remain:
                break
            minColor = min(remain)
            maxColor = max(remain)
            pColor1, pColor2 = self.sortColors(colors, minColor, maxColor, start, end)
            start = pColor1
            end = pColor2
            count += 2
            remain = colors[start:end+1] #start和end都应该包含在下一次里面
        return colors

    def sortColors(self, colors, color1, color2, start, end):
        # left和right表示的是两个颜色'下一个'的位置'
        left = start
        right = end
        index = start
        while index <= right:
            if colors[index] == color1:
                colors = self.swap(colors, left, index)
                left += 1
                index += 1
            elif colors[index] == color2:
                colors = self.swap(colors, right, index)
                right -= 1
            else:
                index += 1

        return left, right


    def swap(self, nums, i, j):
        tmp = nums[j]
        nums[j] = nums[i]
        nums[i] = tmp
        return nums

Reference

• sort colors ii