백준 2213 - 트리의 독립집합 (파이썬)

각각의 정점들에 대하여 자신을 포함했을때의 최댓값과 포함하지 않았을때의 최댓값을 동적 프로그래밍을 사용하여 풀었습니다.

import sys

sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline

n = int(input())
values = [0] + list(map(int, input().split()))
tree = [[] for _ in range(n + 1)]
dp = [[0, 0] for _ in range(n + 1)]
num = [[[], []] for _ in range(n + 1)]

for _ in range(n - 1):
    U, V = map(int, input().split())
    tree[U].append(V)
    tree[V].append(U)

visited = [0] * (n + 1)

def dfs(now):
    visited[now] = 1
    dp[now][0] = values[now]
    num[now][0].append(now)
    
    for next in tree[now]:
        if visited[next] == 0:
            dfs(next)
            dp[now][0] += dp[next][1]
            
            for i in num[next][1]:
                num[now][0].append(i)
                
            if dp[next][0] <= dp[next][1]:
                dp[now][1] += dp[next][1]
                
                for i in num[next][1]:
                    num[now][1].append(i)
            else:
                dp[now][1] += dp[next][0]
                
                for i in num[next][0]:
                    num[now][1].append(i)

dfs(1)

if dp[1][0] >= dp[1][1]:
    i = 0
else:
    i = 1

print(dp[1][i])

tmp = num[1][i]
tmp.sort()
print(*tmp)