데코레이터를 통한 memoization

파이썬 데코레이터를 통한 memoization

피보나치 수열은 재귀 알고리듬의 대표적인 문제인데, 간단히 memoization을 통해서 성능을 개선하는 방법을 찾아보자. 테스트는 파이썬 3.4에서 진행했다.

def fib(n):
    if n in (1, 2):
        return 1
    return fib(n - 1) + fib(n - 2)

위 코드는 피보나치 수열의 정의를 그대로 따르는 코드이다. 이 코드를 사용하여 피보나치 수열의 64번째 항을 구하면… 얼마나 걸릴지는 모르겠다. 실행해놓고 이 글을 다 쓸 때까지도 결과가 나오지 않았다. fib(64)는 재귀 호출을 18,446,744,073,709,551,615 회나 해야하기 때문이다. (\( 2^{64} – 1 \)번을 계산해야 한다) 36번째 항을 계산하는데만해도 14초가 넘게 걸렸다. 이 알고리듬은 매우 간결하고 깔끔한 코드로 보이지만, 불필요한 중복 계산을 너무 많이 한다. 64번째 항을 구하기 위해서는 63번째와 62번째 항을 구하는데 그럼 62번째는 2번 계산된다. 즉 1에 가까워질수록 중복 계산횟수가 기하급수적으로 늘어나는 것이다.

다음 코드는 불필요한 계산을 피하기 위해서 한 번 계산한 결과를 사전에 저장했다가, 계산 결과가 있으면 바로 꺼내서 보여주는 식이다. 이러한 성능 개선 방식을 캐싱 또는 메모이제이션이라 한다.

memo = dict()
def fib2(n):
    if n in memo:
        return memo[n]
    if n in (1, 2):
        memo[n] = 1
        return 1
    result = fib2(n - 1) + fib2(n - 2)
    memo[n] = result
    return result

한 번 계산한 결과를 재계산 할 필요가 없고, 그것이 재귀 호출을 막는 방식이다보니 64번만 계산하면 된다. 다만 이 방식을 사용하는데에는 전역변수를 하나 더 사용해야한다는 불편함이 있다. 만약 이런 비슷한 함수가 여러 종류가 있다면 각 함수마다 따로 따로 중간값을 저장할 공간을 관리해줘야 하는 불편함이 따른다. 이 문제를 좀 더 우아하게 해결할 방법이 없을까?

파이썬의 데코레이터를 보면 함수를 인자로 받아서 함수를 리턴하는 함수를 작성할 수 있다. 이 때 만약 함수 내부의 함수가 부모함수의 스코프에서 특정 값을 캡쳐할 수 있다면 어떨까? 다음 코드를 보자.

def memoize(func):
    tempMemo = dict()
    def wrapped(n):
        if n in tempMemo:
            return tempMemo[n]
        result = func(n)
        tempMemo[n] = result
        return result
    return wrapped

함수 memoize는 함수를 인자로 받고, 다른 함수를 리턴한다. 리턴되는 함수는 임시 사전에 계산 결과를 기록하고, 계산한 이력이 있으면 재계산 없이 해당 내용을 그대로 리턴해준다. 이 때 이 계산 결과를 캐시하는 객체는 부모함수인 memoize의 스코프에 있는 객체가 된다.

함수 memoize는 호출될 때마다 새로운 함수를 리턴하게 되는데, 이 리턴되는 함수는 자신이 리턴되는 시점에 캡쳐한 tempMemo 객체에 대한 참조를 가지고 있다. 따라서 부모함수인 memoize의 사이클은 종료되지만 해당 객체는 파괴되지 않고 계속 유지된다. 이 말은 이 함수를 계속 사용하는 한, 한 번 계산된 결과는 계속해서 캐싱된다는 뜻이다. 이제 다음과 같이 원함수를 바꿔치기하여 계산해보면

fib = memoize(fib)
fib(64)

순식간에 계산된다. 위 함수는 데코레이터의 문법을 그대로 따르고 있어서 코드 작성시에 다음과 같이 쓸 수 있다.

@memoize
def fib(n):
    if n in (1, 2):
        return 1
    return fib(n - 1) + fib(n - 2)

이와 같은 식으로 사용하면 그 외의 함수에서도 동일한 규칙을 적용할 수 있다. 이 때 제약이라함은, 1) 함수는 1개의 인자만 받으며, 2) 그 인자는 사전의 키가 될 수 있는 객체여야 한다. (리스트는 될 수 없다.)