[BOJ] 14267 회사 문화 1 (Python3)

문제

영선회사에는 매우 좋은 문화가 있는데, 바로 상사가 직속 부하를 칭찬하면 그 부하가 부하의 직속 부하를 연쇄적으로 칭찬하는 내리 칭찬이 있다. 즉, 상사가 한 직속 부하를 칭찬하면 그 부하의 모든 부하들이 칭찬을 받는다.

모든 칭찬에는 칭찬의 정도를 의미하는 수치가 있는데, 이 수치 또한 부하들에게 똑같이 칭찬 받는다.

직속 상사와 직속 부하관계에 대해 주어지고, 칭찬에 대한 정보가 주어질 때, 각자 얼마의 칭찬을 받았는지 출력하시오,

입력

첫째 줄에는 회사의 직원 수 n명, 최초의 칭찬의 횟수 m이 주어진다. 직원은 1번부터 n번까지 번호가 매겨져 있다. (2 ≤ n, m ≤ 100,000)

둘째 줄에는 직원 n명의 직속 상사의 번호가 주어진다. 직속 상사의 번호는 자신의 번호보다 작으며, 최종적으로 1번이 사장이다. 1번의 경우, 상사가 없으므로 -1이 입력된다.

다음 m줄에는 직속 상사로부터 칭찬을 받은 직원 번호 i, 칭찬의 수치 w가 주어진다. (2 ≤ i ≤ n, 1 ≤ w ≤ 1,000)

사장은 상사가 없으므로 칭찬을 받지 않는다.

출력

1번부터 n번의 직원까지 칭찬을 받은 정도를 출력하시오.

시간초과 코드

n, m = map(int, input().split())

tree = [[] for _ in range(n + 1)]

infos = list(map(int, input().split()))

valueDic = dict()

for i in range(1, n + 1):
    valueDic[i] = 0
    
for i, info in enumerate(infos):
    if info != -1:
        tree[info].append(i + 1)
print(tree)

def dfs(start, value):
    for g in tree[start]:
        valueDic[g] += value
        dfs(g, value)
        
for _ in range(m):
    i, w = map(int, input().split())
    valueDic[i] += w
dfs(i, w)

for k, v in valueDic.items():
    print(v, end = ' ')
    

 5 3
 -1 1 2 3 4


[[], [2], [3], [4], [5], []]


 2 2
 3 4
 5 6


0 2 6 6 12 

단순하게 dfs를 적용하니까 시간초과가 발생했다.

그래서 재귀를 줄일 방법을 고민하다가, 기존에는 칭찬을 받을때마다 dfs를 돌리며 트리를 내렸다면 이번에는 칭찬을 먼저 받아놓고 1회의 트리 순회로 끝내보는 것으로 코드를 변경하기로 했다. 솔직히 합연산이라 무시해도 된다고 생각했는데, 재귀의 경우에는 합연산도 무시하면 안되겠다.

추가로 recursionError가 발생해서 setresursionlimit을 더 높게 설정해주었다.

정답 코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(1000000)
n, m = map(int, input().split())

tree = [[] for _ in range(n + 1)]

infos = list(map(int, input().split()))

valueDic = dict()

for i in range(1, n + 1):
    valueDic[i] = 0
    
for i, info in enumerate(infos):
    if info != -1:
        tree[info].append(i + 1)
        
for _ in range(m):
    i, w = map(int, input().split())
    valueDic[i] += w
    
def dfs(start):
    value = valueDic[start]
    for g in tree[start]:
        valueDic[g] += value
        dfs(g)
        
dfs(1)

for k, v in valueDic.items():
    print(v, end = ' ')
    

 5 3
 -1 1 2 3 4


[[], [2], [3], [4], [5], []]


 2 2
 3 4
 5 6


0 2 6 6 12