[BOJ] 2213 - 트리의 독립집합 (Python)

 

 

아이디어

 

1949번 우수 마을처럼 DFS와 DP로 가중치의 최댓값을 구하고, 추가로 그 정점 번호들을 저장한다.

 

 

풀이

 

가중치를 리스트 w에 저장하고, 정점 간의 연결을 2차원 리스트 graph에 저장한다.

dp[i][j]는 정점 i를 루트 노드로 하여 해당 정점을 방문했을 때(j = 1)와 방문하지 않았을 때(j = 0)의 최댓값을 저장한다.

S[j]는 해당 정점을 방문했을 때(j = 1)와 방문하지 않았을 때(j = 0)의 최대 독립 집합을 저장한다.

 

함수 dfs는 DFS로 정점을 차례로 방문하여 리스트 visited에 방문 표시 후 리프 노드부터 dp와 최대 독립 집합을 구한다.

우수 마을과 달리 함수 종료 시 현재 정점을 방문했을 때와 방문하지 않았을 때의 최대 독립 집합을 반환한다.

 

현재 정점을 포함하지 않을 때는 자식 정점의 포함 여부와 상관없이 더 큰 최댓값을 더하고 그 집합을 S[0]에 추가하고,

현재 정점을 포함할 때는 자식 정점이 우수 마일이 아닐 때의 최댓값을 더하고 그 집합을 S[1]에 더한다.

 

정점 1을 루트 노드로 함수 dfs를 호출한 후, 가중치 최대 합과 그 최대 독립 집합을 출력한다.

 

import sys

sys.setrecursionlimit(10002)
input = sys.stdin.readline

def dfs(idx):
    dp[idx][1] = w[idx]
    visited[idx] = 1
    S = [set(), set({idx})]
    for i in graph[idx]:
        if not visited[i]:
            temp = dfs(i)
            dp[idx][0] += max(dp[i])
            dp[idx][1] += dp[i][0]
            if dp[i][0] > dp[i][1]:
                S[0] = S[0].union(temp[0])
            else:
                S[0] = S[0].union(temp[1])
            S[1] = S[1].union(temp[0])
    return S[0], S[1]

n = int(input())
w = [0]+list(map(int, input().split()))
graph = [[] for _ in range(n+1)]
dp = [[0]*2 for _ in range(n+1)]
visited = [0]*(n+1)
for _ in range(n-1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)
res = dfs(1)
if dp[1][0] > dp[1][1]:
    print(dp[1][0])
    print(*sorted(res[0]))
else:
    print(dp[1][1])
    print(*sorted(res[1]))

 

 

집합을 정렬하지 않고 그대로 언패킹하여 출력하는 코드도 통과할 수 있었다.

# 집합을 정렬하지 않고 언패킹하는 코드
res = dfs(1)
if dp[1][0] > dp[1][1]:
    print(dp[1][0])
    print(*res[0])
else:
    print(dp[1][1])
    print(*res[1])