[알고리즘] 색 정렬

색 정렬

책 풀이

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        mid = 1
        i, j = 0, 0
        k = len(nums)
        while j < k:
            if nums[j] < mid:
                nums[i], nums[j] = nums[j], nums[i]
                i += 1
                j += 1
            elif nums[j] > mid:
                k -= 1
                nums[k], nums[j] = nums[j], nums[k]
            else:
                j += 1

이 문제는 다익스트라가 1976년에 제안한 네덜란드 국기 문제와 동일한 문제로 퀵 정렬의 개선 아이디어와도 관련이 깊다.

i, k를 양쪽 포인터로 두고, j가 이동하면서 mid 값을 기준으로 스왑하는 형태다.