[백준] 1135번 뉴스전하기 - 파이썬/트리

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

import sys
sys.setrecursionlimit(10**6)

N = int(input())

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

graph = [[] for _ in range(N)]
visited = [False] *(N)
dp = [0] *(N)

for i in range(N):
    if i==0:
        continue
    graph[i].append(S[i])
    graph[S[i]].append(i)


def dfs(x):
    visited[x]=True
    elem = [] # DP값을 담을 리스트
    for i in graph[x]:
        if not visited[i]:
            dfs(i)
            elem.append(dp[i])
    if not elem: # 리프노드면 끝냄
        return
    elem.sort(reverse=True)
    max_num = 0
    for i in range(len(elem)):
        check = elem[i] + (i+1)
        if check>max_num:
            max_num=check

    dp[x] = max_num

dfs(0)
print(max(dp))

트리+DP 유형의 문제이다.

DP 그래프를 0으로 초기화하고 아래부터 접근했다.

• 노드 2

2 노드는 3,4 두개의 리프노드를 가진다.

DP[3], DP[4] 는 리프노드이므로 0이다.

그렇다면 2에서 자식노드들의 DP값을 리스트로 정리해보면 [0,0] 이고, 자식 노드가 2개이므로 전화를 돌리는 시간은 최대 2분이다. 이것을 리스트로 나타내면 [1,2] 로 표현할 수 있다.
두 리스트에서 인덱스가 같은 자리의 합을 구하면 [1,2]이고 여기의 최대값이 DP[2] 값이 된다.

• 노드 0

노드 0에서 보면 자식 노드가 1,2이고, 각각의 DP값을 표현한 리스트는 [DP[1],DP[2]] = [0,2] 이다.

자식노드간의 DP값이 다르다면 DP값이 큰 노드를 먼저 전화해야하므로, 리스트를 내림차순으로 정렬해준다.

그러고나서 자식노드가 2개이므로 [1,2] 의 시간리스트와 더해준다.

  [2,0] # DP값 리스트
  [1,2] # 시간 리스트
+ -------
  [3,2]
  
max([3,2]) = 3

합쳐서 나온 리스트인 [3,2]의 최대값이 DP[0]이 된다.

이러한 방식으로 DP[0]의 값을 구하면 정답이 된다.