import sys
sys.setrecursionlimit(20000)
def input():
return sys.stdin.readline().rstrip()
def trace(cur_node,prev_check):
visited[cur_node] = False
flag = 0
if not prev_check:
if dp[cur_node][0] < dp[cur_node][1]:
result.append(cur_node)
flag = 1
for next_node in graph[cur_node]:
if visited[next_node]:
trace(next_node,flag)
def dfs(node):
visited[node] = False
child_node = []
for next_node in graph[node]:
if visited[next_node]:
child_node.append(next_node)
if not len(child_node):
dp[node][1] = arr[node]
return
else:
dp[node][1] += arr[node]
for child in child_node:
dfs(child)
dp[node][0] += max(dp[child][1],dp[child][0])
dp[node][1] += dp[child][0]
N = int(input())
arr = [0] + list(map(int,input().split()))
# 1이 선택
# 0이 선택 x
dp = [[0 for _ in range(2)] for _ in range(N+1)]
graph = [[] for _ in range(N+1)]
visited = [True for _ in range(N+1)]
for k in range(N-1):
x,y = map(int,input().split())
graph[x].append(y)
graph[y].append(x)
dfs(1)
print(max(dp[1]))
result = []
visited = [True for _ in range(N+1)]
trace(1,0)
result.sort()
print(*result)
어려웠던 문제였다. 트리dp를 활용해서 최적의 경우를 찾는 것은 많았지만, 그 최적을 만족하는 것을 trace를 하는 문제는 처음이라 푸는데 오래걸렸다.
dp라는 테이블을 만들고 (N+1)*2의 배열을 만들어 두었다.
여기서 1은 해당 노드를 선택했을때의 최대값
0은 해당 노드를 선택하지 않았을때의 최대값이다.
만약 리프노드이면 자신을 선택하는 조건이외에는 없을것이다.
그러므로 dp[leef_node][1] = arr[leef_node]를 넣어준다.
리프노드를 제외한 나머지 노드들은 0번 인덱스와 1번인덱스의 최대값을 구해야한다.
현재 노드를 선택했을때 즉 1번인덱스에는
현재 노드를 cur_node라고 했을때 arr[cur_node]를 더해줘야할것이다.
그리고 현재 노드를 선택했으므로, 자식노드들은 선택을 하지 못한다.
그러므로 dp[cur_node][1] += dp[child_node][0] 와 같이
자식 노드들의 선택하지않았을때의 값들을 전부 더해줘야한다.
그리고 현재 노드를 선택하지 않았을때에는
자식 노드를 선택하던지, 선택하지 않는것은 마음대로 할 수 있다.
그러므로 둘중 더 큰 값을 더해주면 된다.
이렇게 dp에 값을 누적시켜서 구하면
우리는 1번노드로 시작했기 때문에
1번노드의 dp에서 최대값을 구할 수 있다.
그러면 우리는 이 dp에서 어떤 노드들이 선택됬는지를 찾아봐야한다.
우리는 이전 노드에서 선택했는지 안했는지에 따라, 현재 노드를 선택할수 있는지 아닌지가 결정이 된다.
그래서 우리는 prev_check를 통해 문제를 해결했다. 위와 동일하게 하기 위해서 0일때에는 부모노드를 선택하지 않았다.
1일때에는 부모노드를 선택했다이다.
만약 부모노드를 선택하지 않았을때에는, 우리는 지금 있는 현재노드를 선택할지 안할지를 결정지을수 있다.
그 결정방법은 dp에 저장된 값을 비교를 해서, 우리가 선택하지 않았을때보다 선택했을때의 값이 더 크면,
그때 현재노드를 선택을 하고, flag를 1로 바꿔준다. 현재노드를 선택을 했으므로, result에 추가 시켜준다.
그리고 재귀를 통해 모든 노드들에 대해 탐색을 해주면 된다.
이 문제는 최대값을 구하는 것 뿐만 아니라, 그 경우의 수가 만들어지는 것을 찾는데 더 어려움을 겪었다.
트리 dp와 관련된 문제들을 좀 더 연습해야겠다.