[백준] 1949 - 우수 마을 [Python(파이썬)]

문제

www.acmicpc.net/problem/1949

1949번: 우수 마을

풀이

n번 마을을 우수 마을로 선정하면 인접한 마을은 모두 우수 마을로 선정할 수 없다.

또한, n번 마을을 우수 마을로 선정하지 않은 상태에서 인접한 마을을 모두 우수 마을로 선정하는 것이 최적해를 보장하지는 않는다. 

 

이 문제는 서브트리의 가중치의 합을 구하는 문제이다. 이를 dfs와 동적계획법으로 풀 수 있다.

dp[n][1]은 n번 마을을 우수마을로 선정했을 경우에 우수 마을의 주민 수의 총합이다.

dp[n][0]은 n번 마을을 우수마을로 선정하지 않았을 경우에 우수 마을의 주민 수의 총합이다.

 

• dfs) u번 마을을 방문하지 않았다면 u번 마을을 방문한다.
• n번 마을을 선정했을 경우 인접한 u번 마을을 선정할 수 없다.
  • dp[n][1] += dp[u][0]
• n번 마을을 선정하지 않은 경우 인접한 u번 마을을 선정하거나 선정하지 않는 두 가지 경우를 모두 고려해야 한다. 
  • dp[n][0] += max(dp[u][0], dp[u][1])

결국 재귀적으로 서브트리의 가중치의 합이 구해지게 되며, 이를 통해 답을 구할 수 있다.

 

코드

import sys, collections
sys.setrecursionlimit(10**6)

def dfs(cur):
    visited[cur] = 1
    for u in g[cur]:
        if not visited[u]:
            dfs(u)
            dp[cur][1] += dp[u][0] # 현재 마을을 우수마을로 선정
            dp[cur][0] += max(dp[u][0], dp[u][1]) # 현재 마을을 우수마을로 선정 X

n = int(sys.stdin.readline().strip())
cost = [0] + [int(x) for x in sys.stdin.readline().split()]

visited = [0 for _ in range(n+1)]
dp = [[0, cost[i]]*2 for i in range(n+1)] # dp[i][0] = i마을을 선정X, dp[i][1] = i마을을 선정O
g = collections.defaultdict(list)

for _ in range(n-1):
    v, u = map(int, sys.stdin.readline().split())
    g[v].append(u)
    g[u].append(v)

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