[Python] 백준 / gold / 2213 : 트리의 독립집합

문제 링크 : https://www.acmicpc.net/problem/2213

트리 DP 문제. 트리랑 각 정점의 가중치가 양의 정수로 주어져 있을 때, 최대 독립 집합을 구하는 문제다.

dp[n][0] = (n을 루트로 하는 서브트리의 최대 독립집합 크기. n 포함 x)
dp[n][1] = (n을 루트로 하는 서브트리의 최대 독립집합 크기. n 포함 o)

path[n][0] = (n을 루트로 하는 서브트리의 독립집합 원소. n 포함 x)
path[n][1] = (n을 루트로 하는 서브트리의 독립집합 원소. n 포함 x)

전에 풀었던 우수 마을(gold2) 문제와 상당히 비슷하다. 이 문제(gold1)는 집합에 속한 노드들도 출력해야되기 때문에 난이도가 조금 더 높았다.
path 가 필요 없었다면, 우수마을 문제와 마찬가지로 리턴값이 없는 dfs 를 설계했어도 됐을거같다.

내가 기억할 것
( dp[node][0] / dp[node][1] ) 로 케이스를 분할하는 트리 dp 문제라고해서, dfs(node, k) 와 같이 매개변수를 두 개 할당해줄 필요가 없었다.

파이썬 코드

import sys
N = int(sys.stdin.readline())
tree = [[] for _ in range(N+1)]
dp = [[0, 0] for _ in range(N+1)]
path = [[[] for _ in range(2)] for _ in range(N+1)]
W = [0] + list(map(int, sys.stdin.readline().split()))
visit = [False] * (N+1)

for _ in range(N-1):
    A, B = map(int, sys.stdin.readline().split())
    tree[A].append(B)
    tree[B].append(A)


def dfs(node):
    visit[node] = True
    dp[node][1] += W[node]
    path[node][1].append(node)

    for x in tree[node]:
        if not visit[x]:
            result = dfs(x)
            dp[node][0] += max(dp[x][0], dp[x][1])
            dp[node][1] += dp[x][0]

            path[node][1] += result[0]
            if dp[x][0] > dp[x][1]:
                path[node][0] += result[0]
            else:
                path[node][0] += result[1]

    return path[node]


p = dfs(1)
if dp[1][0] > dp[1][1]:
    print(dp[1][0])
    p[0].sort()
    for i in p[0]:
        print(i, end=' ')
else:
    print(dp[1][1])
    p[1].sort()
    for i in p[1]:
        print(i, end=' ')