[리트코드 (LeetCode)] Sort Colors (Python)

Sort Colors

• https://leetcode.com/explore/interview/card/top-interview-questions-medium/110/sorting-and-searching/798/
• 분류: Top Interview Questions - Medium - Sorting and Searching

Explore - LeetCode

문제

• 주어진 배열 nums는 n개의 빨간색/하양색/파란색으로 색칠된 오브젝트로 구성되어 있다. 같은 색상이 인접하고 빨-하-파 순이 되도록 in-place로 정렬하시오. 라이브러리의 sort 함수를 쓰지 않고 해결하시오.
• type nums: List[int] / rtype: None Do not return anything, modify nums in-place instead.

예시

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]

조건

• n은 1~300, nums[i]는 0, 1, 2

풀이

이 문제는 Dutch national flag problem(Dutch partitioning problem) 이다.

Dutch national flag problem - Wikipedia

• Dutch national flag problem: 빨강, 하양, 파란색의 공이 있을 때 같은 공끼리 연속하게/정해진 색의 순서대로 나열(정렬)하는 문제
• 주어진 배열을 in-place 방식으로 non-decreasing order(오름차순)으로 정렬하면 됨

1. 계수 정렬 (counting sort)

입력 배열을 순회하여 0, 1, 2의 숫자를 센 후 각 숫자만큼 새 배열에 추가하여 nums에 복사한 후 반환

from collections import Counter
class Solution(object):
    def sortColors(self, nums):
        c = dict(Counter(nums))
        ans = []
        if 0 in c:
            ans += [0] * c[0]
        if 1 in c:
            ans += [1] * c[1]
        if 2 in c:
            ans += [2] * c[2]
        nums[:] = ans

map을 이용해 단축한 코드

from collections import Counter 
class Solution(object):
    def sortColors(self, nums):
        c = dict(Counter(nums))
        tmp = list(map(lambda x: [x]*c[x], c))
        ans = []
        for i in tmp:
            ans += i
        nums[:] = ans

Time: O(n) / Space: O(n)

Runtime: 16 ms (beats 87.48%) / Memory Usage: 13.5 MB (beats 44.76%)

• Counter 객체를 dict로 형변환하지 않아도 잘 동작
• map 객체를 list로 형변환하지 않아도 잘 동작

2. 3 Pointers In-place solution

아이디어

배열을 red, white, blue, unclassified의 네 그룹으로 분류 / red, white, blue 세 개의 포인터를 두고 white 포인터가 blue 포인터보다 작을 때까지 iterate

red 포인터는 마지막으로 red가 놓인 인덱스 +1을 가리킴

white 포인터는 배열을 순회하여 정렬함 (마지막으로 white가 놓인 인덱스 +1을 가리킴)

blue 포인터는 last unclassified element를 가리킴 (그러므로 맨 뒤에서부터 하나씩 감소)

white가 blue를 넘어서면 정렬이 끝났다는 뜻이므로 종료

 

실제 코드 구조

(1) white 포인터가 red를 가리키면 (nums[white] == 0) white와 red 포인터가 가리키는 두 값을 swap하고 두 포인터를 하나씩 전진

(2) white 포인터가 white를 가리키면 올바른 위치이므로 swap없이 한 칸 전진

(3) white 포인터가 blue를 가리키면 가장 마지막의 분류되지 않은(unclassified) 원소와 swap

def sortColors(self, nums):
    red, white, blue = 0, 0, len(nums)-1
    
    while white <= blue:
        if nums[white] == 0:
            nums[red], nums[white] = nums[white], nums[red]
            white += 1
            red += 1
        elif nums[white] == 1:
            white += 1
        else:
            nums[white], nums[blue] = nums[blue], nums[white]
            blue -= 1

Time: O(n) / Space: O(1)

Runtime: 14 ms (beats 92.97%) / Memory Usage: 13.5 MB

 

이 방법이 처음엔 잘 와 닿지 않는데 테스트 케이스 하나 적어두고 코드 따라가 보면 금방 이해된다. 계수 정렬이 항상 유효하지는 않고 또 공간복잡도도 더 높으므로 in-place 정렬 방법도 공부해보는 것을 추천한다.