Python 재귀 없는 DFS 구현 (리턴 값 있어도 OK)

RecursionError: maximum recursion depth exceeded in comparison

아무리 sys.setrecursionlimit을 조절해도 MemoryLimit 메모리 이슈가 발생해서 풀리지 않는 문제가 있었다.

이걸 쓰고 바로 해결했다.

 

Codeforces에서 submission된 코드의 언어를 선택할 수 있다는 것을 알게 된 이후, 다른 참가자들의 파이썬 풀이를 보다 알게 된 기법이다. (https://codeforces.com/contest/1746/submission/181101450)

 

정말 간단한데, 다음과 같이 사용하면 된다.

 

먼저 bootstrap 함수를 정의한다.

from types import GeneratorType
 
def bootstrap(f, stack=[]):
    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

그리고 bootstrap을 dfs 함수의 데코레이터로 활용한다.

이때 return을 모두 yield로 바꿔주기만 하면 끝이다.

 

간단 예제)

def dfs(n):
    if n == 0:
        return 0
    return 1 + dfs(n - 1)

변경 후

@bootstrap
def dfs(n):
    if n == 0:
        yield 0
    yield 1 + (yield dfs(n - 1))

 

실 사용 예제)

@bootstrap
def dfs(u, adj, ss, ks, cum_scores):
    global ans
    k = ks[u]
    nn = len(adj[u])
 
    ans += ss[u] * k
    
    if nn == 0: # leaf
        cum_scores[u] = ss[u]
        yield cum_scores[u]
 
    scores = []
    ks_ = [k // nn for _ in range(nn)]
 
    for v, k_ in zip(adj[u], ks_):
        ks[v] = k_
        score = yield dfs(v, adj, ss, ks, cum_scores)
        scores.append(score)
    
    vals = sorted(scores, reverse=True)
    for val in vals[:k % nn]:
        ans += val
 
    cum_scores[u] = ss[u] + vals[k % nn]
    yield cum_scores[u]

이게 잘 작동하는 원리는, 콜스택을 제네레이터 함수를 스택에 쌓는 방식으로 mimic했기 때문이다.

상당히 똑똑한 방식인 것 같다.

 

앞으로도 유용한 정보로 찾아오도록 하겠다!