題目出處

143 · Sort Colors II

難度

medium

個人範例程式碼

from typing import (
    List,
)

class Solution:
    """
    @param colors: A list of integer
    @param k: An integer
    @return: nothing
    """
    def sort_colors2(self, colors: List[int], k: int):
        # write your code here
        # sort 1~k
        self.partition(colors, 1, k, 0, len(colors)-1)
        return

    def partition(self, colors, color_from, color_to, start, end):
        # end
        if color_from == color_to or start == end:
            return 

        left, right = start, end
        pivot = (color_from + color_to) // 2

        # define
        while(left <= right):
            while(left <= right and colors[left] <= pivot): # if same, left part
                left += 1
            while(left <= right and colors[right] > pivot):
                right -= 1
            if(left <= right):
                colors[left], colors[right] = colors[right], colors[left] # careful infinite loop: left, right still can move

        else: # split
            # print(colors, start, right, left, end)
            self.partition(colors, color_from, pivot, start, right)
            self.partition(colors, pivot+1, color_to, left, end)

算法說明

此題是考 partition 的綜合題。

這題會經常錯在遞迴判斷式與邊界條件，導致無法 AC。

遞迴判斷式

遞迴的定義

# define
while(left <= right):
    while(left <= right and colors[left] <= pivot): # if same, left part
        left += 1
    while(left <= right and colors[right] > pivot):
        right -= 1
    if(left <= right):
        colors[left], colors[right] = colors[right], colors[left] # careful infinite loop: left, right still can move

遞迴的拆解

這邊要處理的是邊界問題，往前我們還需要多注意 while(left <= right)，
這會導致我們是「交錯」才結束，

結束時的狀態為 「start < right < left < end」，由於與 pivot 相等的在兩側都有可能出現，
拆解的搜尋條件為

• start ~ right
• left ~ end

self.partition(colors, color_from, pivot, start, right)
self.partition(colors, pivot+1, color_to, left, end)

遞迴的中止

如果「color_from color_to」 or 「start end」，return

• 「color_from color_to」：已經同顏色，不必再分
• 「start end」：同位置，不必再分

input handling

這題沒有特別說要怎麼處理特殊的 input，預設的 input 都給正確的。

Boundary conditions

上方已經說明完囉

另解 – 更換中止條件

from typing import (
    List,
)

class Solution:
    """
    @param colors: A list of integer
    @param k: An integer
    @return: nothing
    """
    def sort_colors2(self, colors: List[int], k: int):
        # write your code here
        # sort 1~k
        self.partition(colors, 1, k, 0, len(colors)-1)
        return

    def partition(self, colors, color_from, color_to, start, end):
        # end
        if color_from == color_to:
            return 

        left, right = start, end
        pivot = (color_from + color_to) // 2

        # define
        while(left <= right):
            while(left <= right and colors[left] <= pivot): # if same, left part
                left += 1
            while(left <= right and colors[right] > pivot):
                right -= 1
            if(left <= right):
                colors[left], colors[right] = colors[right], colors[left] # careful infinite loop: left, right still can move

        else: # split
            # print(colors, start, right, left, end)
            if start < right: # still need sorting
                self.partition(colors, color_from, pivot, start, right)
            if left < end: # still need sorting
                self.partition(colors, pivot+1, color_to, left, end)

更換部分

• 當顏色一樣的時候，中止

# end
if color_from == color_to:
    return 

• 限制範圍才持續往下做 partition (其實也等同於 start end)

if start < right: # still need sorting
    self.partition(colors, color_from, pivot, start, right)
if left < end: # still need sorting
    self.partition(colors, pivot+1, color_to, left, end)

Reference

• 颜色分类 II · Sort Colors II