[Leetcode] 75. Sort Colors

문제 바로가기
pivot을 기준으로 작은 쪽, 큰 쪽으로 나눠 정렬하는 quick sort의 특징을 활용한다.

Time Complexity: O(kn) - k: 색의 종류
Space Complexity: O(1)

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        n = len(nums)
        for color in range(2, 0, -1):
            storeIdx = 0
            for i in range(n):
                if nums[i] < color:
                    nums[storeIdx], nums[i] = nums[i], nums[storeIdx]
                    storeIdx += 1
            n = storeIdx

Quick Sort

Time Complexity: O(nlogn)
Space Complexity: O(1)

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        self.quicksort(nums, 0, len(nums) - 1)
        
    def partition(self, a, left, right, pivotIdx):
        pivotValue = a[pivotIdx]
        a[pivotIdx], a[right] = a[right], a[pivotIdx]  # 피벗을 끝으로 옮겨 놓는다.
        storeIdx = left
        for i in range(left, right):
            if a[i] <= pivotValue:
                a[storeIdx], a[i] = a[i], a[storeIdx]
                storeIdx += 1
        a[right], a[storeIdx] = a[storeIdx], a[right]  # 피벗을 두 리스트 사이에 위치시킨다.
        return storeIdx

    def quicksort(self, a, left, right):
        if right > left:
            pivotNewIdx = self.partition(a, left, right, right)
            self.quicksort(a, left, pivotNewIdx-1)
            self.quicksort(a, pivotNewIdx+1, right)