백준 2213번 : 트리의 독립집합 in Python

2213번: 트리의 독립집합

 

www.acmicpc.net

이 문제는 트리에서 정점의 가중치의 합이 최대인 독립 집합을 찾는 문제이다. 독립 집합이란 정점의 부분 집합으로 집합 내의 모든 정점쌍이 서로 인접하지 않는다는 특징을 갖는다.

만약 이 문제가 일반적인 그래프에 대해 주어지면 NP hard 문제가 된다. 여기서 NP hard 문제는 일반적으로 polynomial time 안에 해결할 수 없다는 것으로 추측되는 문제를 말한다. 즉, input size가 클 경우 문제를 해결하는 것이 매우 오랜 시간이 걸려 사실상 문제 해결이 불가능하다는 의미이다. 

그러나 트리에 대해 최대 독립 집합을 구하는 경우라면 트리의 특수성을 이용해 polynomial time 안에 문제를 해결할 수 있다. 우선 트리 내에서 루트 노드가 될 정점을 임의로 하나 선정하고, 해당 정점과 인접한 정점들을 자식 노드로 갖게 한다. 그 다음 dynamic programming을 이용해 bottom up 방식으로 leaf 노드부터 차례로 올라오면서 최대 독립 집합을 만들어 나간다.

dynamic programming의 원리는 어떤 정점 u에 대해 계산할 때 u를 루트로 하는 서브트리에 대해 u를 포함하는 최대 독립 집합과, u를 포함하지 않는 최대 독립 집합을 만드는 것이다. u를 포함하는 최대 독립 집합의 경우 u와 인접한 정점들을 포함하면 안 되는 조건이 있기 때문에 u의 각 자식 노드를 루트로 하는 서브 트리에서 자식 노드를 포함하지 않는 최대 독립 집합을 u의 최대 독립 집합에 더한다. u를 포함하지 않는 최대 독립 집합의 경우 u와 인접한 정점들을 포함하든 말든 상관 없기 때문에 u의 각 자식 노드를 루트로 하는 서브 트리에서 자식 노드를 포함하는 최대 독립 집합과 자식 노드를 포함하지 않는 최대 독립 집합을 비교해 정점의 가중치의 합이 더 큰 집합을 선택해 u의 최대 독립 집합에 더한다.

base case인 leaf 노드인 경우 해당 노드를 포함한 경우와 포함하지 않는 경우가 각각 노드 하나만 가지는 집합과 공집합이 된다. 따라서 leaf 노드부터 bottom up 방식으로 올라가면서 최대 독립 집합을 구해나가면 마지막에는 루트 노드에 대해 최대 독립 집합을 구하게 되고, 그러면 모든 노드들을 포함하는 트리에 대해 최대 독립 집합을 구한 것과 같다. 이러한 방식은 각 노드에 대해 그와 인접한 정점들을 살펴보는 시간만 소요되므로 O(|V||E|) 시간만에 답을 얻을 수 있어 효율적이다.

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

import sys
input = sys.stdin.readline

n = int(input())
weights = [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)
# DFS 방식으로 구현하기 위해 사용하는 visited 배열
visited = [False for _ in range(n+1)]
# dynamic programming을 위한 행렬로 
# 각 노드마다 해당 노드를 포함한 경우 최대 독립 집합의 가중치의 합과 
# 해당 노드를 포함하지 않은 경우 최대 독립 집합의 가중치의 합을 저장한다.
dp = [[weights[i], 0] for i in range(n+1)]
# dynamic programming을 위한 행렬로 
# 각 노드마다 해당 노드를 포함한 경우 최대 독립 집합과 
# 해당 노드를 포함하지 않은 경우 최대 독립 집합을 저장한다.
indep_set = [[[i], []] for i in range(n+1)]
def DFS(currNode):
    visited[currNode] = True
    for v in graph[currNode]:
        if not visited[v]:
            # 자식 노드를 재귀적으로 호출해서 bottom up 방식을 구현했다.
            DFS(v)
            # 서브트리의 루트 노드를 포함한 경우
            dp[currNode][0] += dp[v][1]
            indep_set[currNode][0] += indep_set[v][1]
            # 서브트리의 루트 노드를 포함하지 않은 경우
            if dp[v][1] > dp[v][0]:
                dp[currNode][1] += dp[v][1]
                indep_set[currNode][1] += indep_set[v][1]
            else:
                dp[currNode][1] += dp[v][0]
                indep_set[currNode][1] += indep_set[v][0]
# 임의의 노드를 루트 노드로 정해도 된다.
DFS(1)
indep_set[1][0].sort()
indep_set[1][1].sort()
# 루트 노드를 포함한 경우와 포함하지 않는 경우 중 더 큰 가중치를 갖는 독립 집합을 고른다.
if dp[1][0] > dp[1][1]:
    print(dp[1][0])
    print(" ".join(map(str, indep_set[1][0])))
else:
    print(dp[1][1])
    print(" ".join(map(str, indep_set[1][1])))