[Python] 백준 15462번: The Bovine Shuffle (백준 Python 런타임 에러 (RecursionError) 해결하기, sys.setrecursionlimit)

문제

https://www.acmicpc.net/problem/15462

15462번: The Bovine Shuffle

소의 기분을 좋게 해주기 위해서 셔플 노래를 틀어주며 자리를 섞는데, 항상 채워져있는 자리의 개수를 찾는 것이 문제이다!

 

설명

그래프라고 생각하고 cycle이 생기는 곳을 찾아 node의 개수를 찾으면 되는 문제이다. 

코드

import sys
sys.setrecursionlimit(1000000)

N = int(sys.stdin.readline())
next_info = [int(x) for x in sys.stdin.readline().split()]  # i+1

visited = [0 for x in range(N)]
depth = [0 for x in range(N)]

ans = 0

def DFS(curr_pos, cnt, trace):
    global ans
    depth[curr_pos] = cnt
    visited[curr_pos] = 1
    trace += [curr_pos]
    next_pos = next_info[curr_pos]-1
    if visited[next_pos] == 0:  # never visited
        DFS(next_pos, cnt+1, trace)
    else:
        if next_pos in trace:
            ans += (cnt+1-depth[next_pos])
    
for i in range(N):
    if visited[i] == 0:
        DFS(i, 1, [])
        
print(ans)

처음에는 온 곳을 다 저장하는 trace라는 변수를 선언하고 DFS를 돌면서 만약 지나온 곳이라면 cycle이 생성되었다고 판단하는 코드를 작성했다. 노드 개수를 찾기 위해서 DFS에서 몇 번째로 방문한 곳인지 카운트하며 depth라는 벡터에 저장했다. 그래서 지나온 곳을 방문했을 때 DFS의 cnt와 이전에 방문했을 때의 cnt를 뺀 값에 + 1을 하면 cycle을 구성하는 노드의 개수를 구할 수 있다.

근데 그냥 코드를 제출하면 런타임 에러 (RecursionError)라는 결과를 얻는다. 이는 Python의 경우 재귀함수 호출 회수를 제한하기 때문인데 sys.setrecursionlimit()을 통해서 값을 재설정하면 해결할 수 있다.

나중에 고민하다가 안 사실이지만 trace라는 변수를 이용해서 지나온 노드를 저장할 필요도 없다. 해당 코드는 아래와 같다.

import sys

N = int(sys.stdin.readline())
adj = [int(x) for x in sys.stdin.readline().split()]  # i+1

visited = [0 for x in range(N)]

ans = 0
start = 0
for i in range(N):
    if visited[i] == 0:
        bs = start;
        while not visited[i]:
            bs += 1
            visited[i] = bs
            next_pos = adj[i] -1
            i = next_pos
        if (visited[next_pos] > start):
            ans += (bs - visited[next_pos] + 1)
        start = bs
    
print(ans)

 

 

결과

trace를 통해 값을 다 저장할 경우 메모리를 엄청 잡아 먹는 것을 알 수 있다.