[Python] 백준 / gold / 1949 : 우수 마을

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

트리 DP 문제. 이런 생각의 흐름으로 문제를 풀었다.

1. 트리 구조고, N = O(10^4) 이기 때문에 우수 마을을 선정하는 경우의 수를 일일히 완전탐색 해본다면 시간초과가 날거라고 생각했다. 그래서 트리 DP 로 풀어야겠다고 생각이 들었다.

2. DP 는 점화식을 세워야되고, 트리 DP 는 점화식의 주제를 서브트리로 잡아야 한다.
   • dp[n][0] = n을 루트로 하는 서브트리에서, n이 우수 마을로 선정되지 않았을 때 최댓값
   • dp[n][1] = n을 루트로 하는 서브트리에서, n이 우수 마을로 선정되었을 때 최댓값

3. 조건 3번 : "우수 마을로 선정되지 못한 마을은 적어도 하나의 우수 마을과는 인접해야 한다" 이 조건이 문제의 난이도를 높인 것 같다. 처음엔 n 과 인접한 모든 children 을 구해서 children 중 어떤 마을을 우수 마을로 설정할 것인지 combinations 를 구해 그 중 최댓값을 구하려고 시도했다.
   그런데 그렇게까지 할 필요는 없었고, max(dp[child][0], dp[child][1]) 을 더해주는게 더 편한 방법이었다.

4. 로직을 짜고 보니, 연속으로 계속 dp[child][0] 을 선택하게 되면 조건 3번을 불만족하게 되는것이 아닌가 생각이 들었는데, 쫌만 더 생각해보면 그럴 일이 아예 일어나질 않는다. max 로 묶여있기 때문이다.

파이썬 코드

import sys
sys.setrecursionlimit(10**6)
N = int(sys.stdin.readline())
people = [0] + list(map(int, sys.stdin.readline().split()))
tree = [[] for _ in range(N+1)]
dp = [[0, 0] for _ in range(N+1)]
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] += people[node]
    for child in tree[node]:
        if not visit[child]:
            dfs(child)
            dp[node][0] += max(dp[child][0], dp[child][1])
            dp[node][1] += dp[child][0]


dfs(1)
print(max(dp[1]))