백준 1949번 : 우수 마을 in Python

1949번: 우수 마을

 

www.acmicpc.net

이 문제는 기존의 트리의 독립집합을 구하는 문제에 '독립집합에 들어가지 않는 정점은 적어도 하나의 독립집합 내 정점과 인접해야 한다'는 추가 조건이 붙은 경우이다.

이 조건을 만족시키기 위해서 기존의 트리의 독립집합을 구하는 알고리즘에 각 정점 u에 대해 u를 루트로 하는 서브트리에서 u를 포함하지 않는 독립집합의 정점들의 가중치 합의 최댓값을 구할 때 u의 자식 노드 중 자식 노드 자신을 포함하는 독립집합을 적어도 한 번 선택했는지를 boolean 값으로 나타내는 배열 near_good을 추가한다.

기존 문제에서는 dynamic programming 과정에서 인자로 주어진 노드가 독립집합에 포함되지 않을 경우 자식 노드가 독립집합에 포함되는지 안 되는지는 상관이 없었다. 그러나 near_good 배열의 자식 노드에 대응하는 값이 거짓, 즉 자식 노드가 그의 어떠한 자식 노드들도 독립집합에 포함하지 않은 경우 현재 노드까지 독립집합에 들어가지 않으면 자식 노드는 모든 인접한 노드들이 독립집합에 포함되지 않게 된다. 이는 추가된 조건을 위배하므로, 이러한 경우 자식 노드가 독립집합에 포함되지 않는 경우가 독립집합의 가중치 합이 더 커도 강제로 자식 노드를 독립집합에 포함시키는 경우를 선택해야 한다.

이 점만 주의하면 나머지 구현은 이전의 트리의 독립집합 문제와 동일하다.

이를 코드로 구현하면 다음과 같다.

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

n = int(input())
population = [0] + list(map(int, input().split()))
graph = [[] for _ in range(n+1)]
for _ in range(n-1):
    u, v = map(int, input().split())
    graph[u].append(v)
    graph[v].append(u)
visited = [False for _ in range(n+1)]
# 독립집합의 가중치의 합의 최댓값을 두 가지 경우로 나눠 저장하는 배열이다.
# 앞의 경우는 해당 노드를 포함하는 독립집합, 뒤의 경우는 해당 노드를 포함하지 않는 독립집합에 대응한다.
dp = [[population[i], 0] for i in range(n+1)]
# 각 노드에 대해 노드를 포함하지 않는 독립집합에 해당 노드의 
# 자식 노드들 중 적어도 하나를 가지고 있는지를 나타내는 배열이다.
near_good = [False for _ in range(n+1)]
def DFS(currNode):
    visited[currNode] = True
    for v in graph[currNode]:
        if not visited[v]:
            DFS(v)
            dp[currNode][0] += dp[v][1]
            # 만약 near_good[v]가 False면 강제로 자식 노드를 독립집합에 포함시켜야 한다. 
            if dp[v][0] >= dp[v][1] or not near_good[v]:
                # 이 경우 자식 노드를 독립집합에 포함시켰으므로 배열의 값을 True로 바꾼다.
                near_good[currNode] = True
                dp[currNode][1] += dp[v][0]
            else:
                dp[currNode][1] += dp[v][1]        
DFS(1)
print(max(dp[1][0], dp[1][1]))