[Python] 백준 1949번 - 우수 마을

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

1949번: 우수 마을

 

스터디 할 때 트리 DP 예제 코드를 받아놔서 그거 보고 좀 고치면서 풀었더니 골드1 치고는 엄청 빨리 풀었다. 사실 두 시간 만에 푼 거라 난이도 떠나서도 내 기준 엄청 빨리 푼 편에 속한다.

 

문제 이해부터 쉽지 않았는데, 문제에서 중요한 바는

• 각 노드에서 나아갈 수 있는 방향은 '우수마을' or 'not 우수마을' 이렇게 두 가지 경우.
• 루트 노드가 우수 마을이면 하위 노드는 우수 마을 일 수 없다.
• 루트 노드가 우수 마을이 아니면, 하위 노드는 우수 마을이거나 아니거나. (둘 중 최댓값....아마도)

이렇게 세 가지인 것 같다.

 

 

dp 문제라서 dp에 저장할 값을 정하는 게 가장 난관이었다.

우수 마을인 경우와 아닌 경우를 구분해주어야 하므로

 

dp[i][0] : 현재(i) 마을이 우수 마을이 아닌 경우

dp[i][1] : 현재(i) 마을이 우수 마을인 경우

 

이렇게 두 개로 인덱스를 나누어서 저장하였다.

 

if not visited[i]: 내부에 dp값을 저장하는 부분에서

dp[start][0]은 현재 start 마을이 우수 마을이 아니므로, 다음 마을은 우수 마을일 수도, 우수 마을이 아닐 수도 있다. 둘 중 더 큰 값을 선택해서 dp[start][0]에 더해주면 된다.

dp[start][1]은 현재 start 마을이 우수 마을이므로, 다음 마을은 우수 마을일 수 없다. 따라서, dp[i][0]을 더해주면 된다.

 

dp를 다 저장한 후 어떤 값을 출력할 지 정하는 것도 헷갈렸다.

처음에는 그리디처럼 가장 인구 수가 많은 마을에서 시작해야하나? 싶어서 고민이 좀 됐었다.

잘 모르겠어서 일단 예제 값을 넣어보고 예제의 결과값이 dp[1]에 들어있길래 dp[1][0]과 dp[1][1] 중 max 값을 출력하도록 해주었다.

사실.. 찍은 거나 마찬가진데 어쨌든 맞았으니까.. 넘어간다! ㅋㅋ

 

제출하고 나서 recursion error 나서 당황했는데 sys.recursionlimit(10**6) 넣어주니까 바로 풀렸다 !!

 

 

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

def village(start):
    visited[start] = True

    for i in tree[start]:
        if not visited[i]:
            village(i)
            dp[start][0] += max(dp[i][0], dp[i][1])
            dp[start][1] += dp[i][0]
    
n = int(input())
people = [0] + list(map(int, input().split()))
tree = [[] for i in range(n+1)]
visited = [0 for i in range(n+1)]
dp = [[0, people[i]] for i in range(n+1)] #[i][0]: 우수마을 x, [i][1]: 우수마을
for i in range(n-1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)
village(1)
print(max(dp[1][0], dp[1][1]))