[코테] 파이썬 시간효율성을 높이는 몇가지 팁 (feat. 백준 1517)

파이썬은 대표적인 스크립트 언어로, 코드가 직관적이고 간결한 반면 실행속도가 느리다는 단점을 가지고 있습니다.

대부분의 코테에서 가장 쓰기 편한대신 시간효율성은 많이 떨어지는 편입니다. 실제로 대부분의 코테 문제들은 C++나 Rust같은 언어로 풀었을때 가장 빠른 결과를 얻을 수 있습니다.

특정문제의 걸린시간 랭킹. C, C++, Rust같은 언어만 보입니다.

물론 대부분의 코테는 문제만 풀면 그만이지만 (최적화에 신경쓸 필요없다는 소리는 아닙니다 ^^;) 파이썬은 너무나도 느린나머지 올바른 Big-O를 가진 알고리즘을 사용하더라도 몇가지 작은 실수에 의해 시간초과가 발생할 수 있습니다.

즉, 코테 언어로 파이썬을 사용하시는분들은 기본적인 최적화에 관해 필수적으로 고민해보셔야 합니다. 난이도가 낮은문제들은 상관없지만, 난이도가 높아지는경우 파이썬코드 최적화가 잘 되어있지 않으면 체감상 30~40% 정도는 올바른 알고리즘에 대해 시간초과가 발생하는것 같습니다.

비록 제가 남에게 뭘 가르치고 할 실력은 아니지만, 코테 문제를 풀면서 알게된 파이썬 최적화 팁들을 몇가지 공유하고자 합니다.

예시코드와 문제점

이번 포스트에선 간단한 문제(백준 1517번) 하나를 예시로 들어 설명하겠습니다.

해당 문제는 요소 $N \leq 500K$ 개를 가지고있는 수열을 버블정렬할때, 몇번의 스왑이 발생하는지 알아내야하는 문제입니다. 일견 간단해보이지만, 시간제한이 1초입니다. 이때 $O(n^2)$인 알고리즘은 보통 100K 정도에서 시간제한에 걸리게 되므로, 버블정렬은 사용할 수 없습니다. $O(n\log{n})$인 알고리즘으로 푸셔야 합니다.

해당 문제를 푸는 알고리즘은 세그먼트 트리와 병합정렬이 있는데, 병합정렬방식이 보통 빨라서 병합정렬 코드로 설명드리겠습니다.

import sys
sys.setrecursionlimit(10**7)

def merge_sort(arr, left=0, right=None):
    global cnt
    if right is None:
        right = len(arr) - 1

    if left == right:
        return [arr[left]] # 문제 1
   
    mid = left + (right - left) // 2
    left_arr = merge_sort(arr, left, mid) # 문제 1
    right_arr = merge_sort(arr, mid+1, right) # 문제 1

    merged_arr = [] # 문제 1
    i, j = 0, 0
    while i < len(left_arr) and j < len(right_arr):
        if left_arr[i] <= right_arr[j]:
            merged_arr.append(left_arr[i]) # 문제 2
            i += 1
        else:
            merged_arr.append(right_arr[j]) # 문제 2
            j += 1
            cnt += len(left_arr[i:]) # ★문제 3
    merged_arr += left_arr[i:] + right_arr[j:] #문제 2
    return merged_arr

N = int(sys.stdin.readline())
cnt = 0
arr = list(map(int, sys.stdin.readline().split()))

merge_sort(arr)
print(cnt)

가장 처음에 작성한 코드이고, 알고리즘 동작은 $O(n\log{n})$입니다. 하지만 백준에서 시간제한에 걸린 코드입니다. 어느 부분에 문제가 있었는지 주석으로 표시해 두었습니다. 그럼 바로 알아보겠습니다.

1. 잦은 리스트 선언은 피하자.

주석의 문제 1번 코드들을 보시면 재귀가 발생할때마다 왼쪽 리스트와 오른쪽 리스트, 그리고 병합된 결과를 반환할 리스트를 선언하고 있습니다. 더이상 나눌 수 없는 경우에도 길이가 1인 리스트를 새로 생성해서 반환하고 있습니다. 이렇게 코드를 작성하면 공간, 시간적인 측면에서 비효율성이 발생합니다. 주로 공간효율성에 큰 영향을 미치긴하지만, 시간효율성에도 영향이 있습니다. 왜그럴까요?

• 가비지 컬렉터(gc)의 비효율성 초래

  파이썬은 레퍼런스 카운팅과 가비지 컬렉터(이하 gc)로 프로그램 메모리를 관리합니다. 여기서 gc는 동작중에 프로그램을 잠시 중지시키기 때문에, gc가 자주 동작할수록 프로그램이 느려집니다. gc는 보통 많은 메모리를 차지하는 컨테이너객체(list, tuple..)를 다수 선언해서 계산하는경우 순환참조의 가능성이 높아져 자주 동작할 가능성이 높습니다.

  기존 코드에서 merge_sort를 실행한 이후 gc가 몇번 동작했는지를 아래와 같이 추적해 보았습니다.

  import gc
  
  arr = [random.randint(0, 100000) for _ in range(100000)]
  gc.set_debug(gc.DEBUG_STATS)
  
  t1 = gc.get_count()
  merge_sort(arr)
  t2 = gc.get_count()
  print('현재 코드 : ', *(b-a for a, b in zip(t1, t2)))

  매 실행시 결과가 달라지지만, 보통 5 ~ 13회 정도 0세대 가비지 컬렉션이 진행되었습니다. 이후 개선된 코드와 비교해볼 예정입니다.
  
  

• 낮은 참조 지역성 (캐시미스 증가)

  보통 컴퓨터는 자주 사용하는 메모리를 캐시로 이동시켜 CPU와 주 기억장치간의 속도 병목을 최소화 하고자 노력합니다. 이때 캐시를 특정 메모리 블럭단위로 가져오기때문에, 계산에 필요한 정보들의 메모리주소가 붙어있으면 시간효율성이 높아지게 됩니다.

  파이썬의 리스트는 내부적으로 c의 동적배열로 구현되어있습니다. 배열이기때문에 한번 선언하면 배열의 근처요소는 근처의 메모리주소를 가지게 되고, 캐시에 배열의 메모리 블럭을 가져왔을때 다음 계산에 필요한 데이터도 같이 포함되어있을 확률이 높습니다. 따라서 참조 지역성이 좋습니다.

  그런데 기존의 arr를 쓰지않고 left_arr = arr[i:j] 이런식으로 새로운 변수에 arr의 i부터 j까지요소를 선언하면 어떻게 될까요? 일견 left_arr가 arr의 i 부터 j까지 주소를 가져오는것으로 생각할 수 있지만, 실제로는 완전히 새로운 메모리영역에 새로운 배열을 선언하고 i부터 j까지 값을 복사해서 가져오게 됩니다. (좀더 엄밀하게는 값에 인스턴스 객체의 메모리 주소를 복사하겠죠? 확실치 않네요..) 따라서 재귀 할때마다 원본 배열의 주소가 아니라 완전히 다른 주소를 가져와 계산하기때문에, 참조 지역성이 낮아집니다. 그럼 캐시미스가 증가하게되고, 빠른 캐시메모리가 아니라 느린 RAM혹은 보조기억장치에서 데이터를 가져오는 경우가 많아집니다.

  즉, 쓰던 리스트를 안쓰고 새롭게 선언해서 계산하면 공간적 참조 지역성 이슈로 시간효율성이 낮아집니다.
  
  

2. append와 extend사용은 필요할때만

바로 위에서 말씀드렸다싶이 파이썬 리스트는 동적배열로 구현되어있습니다. 따라서 배열의 길이가 부족해지면 길이가 더 긴 완전히 새로운 배열을 선언하고 기존 내용을 복사합니다. 얼마나 늘리는진 버전별로 인터프리터별로 다른걸로 알고있습니다만, CPython은 약 1.125배 길이를 늘린 배열을 선언하는것으로 알고있습니다. (이부분 정확하지 않을 수 있습니다.) 어쨋든 새로운배열을 선언하고 복붙하는 작업을 내부적으로 하기때문에, 한번 실행시 $O(n)$의 시간이 소요됩니다. 물론 append할때마다 새로운 배열을 선언하는것은 아니기때문에 append와 extend의 시간복잡도는 분할상환분석상 $O(1)$이지만, 한번도 배열 재선언을 안하는 경우에 비해선 당연히 오래걸릴 수 밖에 없습니다.

또한 배열을 재선언하게되면 새로운 메모리주소를 사용하게됩니다. 따라서 재선언을 자주하게되면 당연히 참조 지역성도 낮아지겠죠?

따라서 리스트의 크기를 미리 알 수 있는경우 처음부터 정해진 크기로 선언하고 순서대로 값을 할당하는 방식이 시간효율성이 좋습니다.

3. 루프 안에서 대용량 리스트의 슬라이싱 혹은 메서드사용은 가급적 회피할것

사실 주석의 문제 3은 문제 1번에 포함될수도 있는내용이지만, 기존코드에서 가장 큰 시간 비효율성을 초래한 부분이라 따로 빼서 설명드립니다.

코드를 보시면 len(left_arr[i:]) 부분은 while루프안에 들어가있습니다. 앞서 설명드렸다싶이 left_arr[i:] 이런식으로 작성하면 새로운 리스트를 선언하는것이고, 내부적으로 새로운 배열을 선언해서 복붙하게 되고, $O(n)$이 소요됩니다. $O(n)$이 while루프안에있고, $N \leq 500K$ 이므로 left_arr의 최대 요소는 250K 입니다. 즉 해당코드는 최악의 경우 한번에 수십만개 요소를 가진 리스트 선언을 수십만번 해야할 수도 있습니다. 물론 그럴일은 거의 없겠지만, 평균적으로 봐도 매우 비효율적인 코드임을 쉽게 알 수 있습니다.

그런데 생각보다 이부분에서 실수가 자주 발생하는게, 리스트는 기본적으로 자신의 길이를 추적하고 있으므로 len(List) 자체는 $O(1)$입니다. 따라서 당황하거나 긴장해서, 혹은 별생각없이 사용하다가 리스트를 선언하는 경우임을 깜빡하는경우가 종종 있습니다.

이런 실수를 줄이려면 어떤경우에 원본인지, 새로운 리스트인지, 어떤 시간복잡도를 가지는지를 항상 정확히 인지하고 코딩을 해야합니다. 근데 그게 힘들면 다른방법을 생각해봐야죠. 예를들어 cnt += len(left_arr[i:]) 는 cnt += mid - i + 1로 immutable한 int만 써서 계산할 수 있습니다.
아직 파이썬의 정확한 동작에 익숙하지 않다면 루프안에서는 가급적 리스트 슬라이싱이나 메서드사용을 회피하고 다른 방식을 생각해보는것이 좋을 수 있습니다.

개선된 코드

앞서 말씀드렸던 문제점들을 개선한 코드입니다.

import sys
sys.setrecursionlimit(10**7)

def merge_sort2(arr, left=0, right=None, merged=None):
    global cnt
    if right is None:
        right = len(arr) - 1
        merged = [None] * len(arr)

    if left >= right:
        return
    
    mid = left + (right - left) // 2
    merge_sort2(arr, left, mid, merged)
    merge_sort2(arr, mid+1, right, merged)

    i, j, k = left, mid+1, left
    while i <= mid and j <= right:
        if arr[i] <= arr[j]:
            merged[k] = arr[i]
            i += 1
        else:
            merged[k] = arr[j]
            j += 1
            cnt += mid - i + 1
        k += 1
    while i <= mid:
        merged[k] = arr[i]
        i += 1
        k += 1
    while j <= right:
        merged[k] = arr[j]
        j += 1
        k += 1
    
    for n in range(left, right+1):
        arr[n] = merged[n]
    return


N = int(sys.stdin.readline())
cnt = 0
arr = list(map(int, sys.stdin.readline().split()))

merge_sort2(arr)
print(cnt)

변경점
1. 새로운 리스트를 선언하는 코드를 전부 없애고 merged하나에 인덱스로만 처리하도록 바꿨습니다.
2. append와 extend, List + List를 전부 없애고 정해진 크기의 merged 리스트에 할당하는 방식으로 바꿨습니다.
3. while루프 안에서 불필요하게 선언되던 리스트를 없애고 간단한 연산으로 바꿨습니다.

이런식으로 작성하면 merged 리스트를 지속적으로 사용하게되고 불필요한 객체 선언도 없으므로 참조 지역성이 좋아졌습니다. append와 리스트 이어붙이기가 없기때문에 배열 재선언도 없습니다. cnt 계산도 간단해졌습니다.

개선된 코드의 gc 동작횟수도 알아보겠습니다.

arr = [random.randint(0, 100000) for _ in range(100000)]
arr2 = arr.copy()

t1 = gc.get_count()
merge_sort(arr)
t2 = gc.get_count()

t3 = gc.get_count()
merge_sort2(arr2)
t4 = gc.get_count()

print('현재 코드 : ', *(b-a for a, b in zip(t1, t2)))
print('개선 코드 : ', *(b-a for a, b in zip(t3, t4)))

현재 코드 : 14 0 0
개선 코드 : 0 0 0

개선 코드는 0 0 0 으로 모든세대에서 가비지 컬렉션이 한번도 실행되지 않았음을 확인할 수 있었습니다.

결과적으로, 개선된 코드는 같은 알고리즘으로 시간초과없이 문제를 푸는데 성공했습니다. C++같은 언어에 비하면 여전히 매우 느리지만, 너무느려서 알고리즘이 맞는데도 틀리는 경우만 아니면 충분하죠.

기타 간단 팁

작성하고보니 오롯이 리스트에 집중한 내용이 되었습니다. 파이썬으로 코테를 풀때 가장 많이 사용하는게 리스트라서 이것의 동작을 이해하는게 많이 중요한것 같습니다.

그래도 아쉬우니 리스트가 아닌 다른 간단팁 몇개 뿌리고 마칩니다.

1. 코드수정없이 PyPy3으로 시도해보기 : 백준같은 경우는 PyPy3을 지원하는데, 이건 컴파일러를 내장하고있어 완전 스크립트인 CPython보다 어느정도 효율성이 더 나오는 경우가 많습니다. Python3로 시도했을때 시간초과났던 코드가 PyPy3으론 될 수도 있습니다. 단, PyPy3은 재귀에 약한편입니다. 고난이도 재귀문제를 PyPy3로 풀면 메모리초과의 위험성이 있습니다.

2. 바뀔일 없는 데이터는 tuple로 관리 : 튜플은 immutable한 컨테이너 객체이므로 일반적으로 리스트보다 요소 접근속도가 빠릅니다. 자르거나 추가하거나, 수정할일이 없는 데이터는 튜플로 관리하면 소폭의 속도상승을 기대해볼 수 있습니다.

3. 딕셔너리(해시테이블) 사용 고려 : 파이썬은 해시테이블에 강점을 가지고 있는 언어 입니다. 해시테이블은 메모리 비효율성이 높은대신 입출력 시간에 상당한 강점을 가지고있습니다. 대부분의 문제는 메모리제한보다는 시간제한이 더 타이트하므로 딕셔너리 사용을 고려하기 좋은 환경입니다. 문제 유형이 해시테이블이 아니더라도 데이터를 임시저장하거나 처리할때 리스트나 튜플대신 딕셔너리를 사용해보세요.