[백준 2213번] 트리의 독립집합 (Python3)

import sys
input = sys.stdin.readline

N = int(input())
weight = [*map(int,input().split())]

graph = [set() for i in range(N)]

for i in range(N-1):
  x,y = map(lambda x:x-1,map(int,input().split()))
  graph[x].add(y)
  graph[y].add(x)

DP = [[0,0] for i in range(N)]
path = [[set(),set([i])] for i in range(N)]

def DFS(now):
  for next in graph[now]:
    graph[next].remove(now)
    DFS(next)
    if DP[next][0] >= DP[next][1]:
      DP[now][0] += DP[next][0]
      path[now][0].update(path[next][0])
    else:
      DP[now][0] += DP[next][1]
      path[now][0].update(path[next][1])
    DP[now][1] += DP[next][0]
    path[now][1].update(path[next][0])
  DP[now][1] += weight[now]

DFS(0)
if DP[0][0] >= DP[0][1]:
  print(DP[0][0])
  print(" ".join(map(str,map(lambda x:x+1,sorted([*path[0][0]])))))
else:
  print(DP[0][1])
  print(" ".join(map(str,map(lambda x:x+1,sorted([*path[0][1]])))))

 

DFS + DP 문제였다.

사회망 서비스 (SNS) 문제와 거의 유사해서 쉽게 해결할 수 있었는데, 경로추적하는게 약간 까다로웠다.

 

알고리즘을 요약하면

1. DP를 2*N짜리 리스트로 만든다. 이때 DP[now][0]은 집합에 node now가 속하지 않는것, DP[now][1]은 속하는 것을 의미한다.

2. 그래프로 연결된 다음 노드로 DFS를 실행한다. 이때 양방향 그래프이므로 이동하려는 next 노드에서 now 노드로의 간선을 제거해주어야 한다.

3. [now][0]은 next노드가 독립집합에 포함될 경우와 포함되지 않을 경우를 비교해서 더 큰값을 가져오고, [now][1]은 포함되지 않을 경우만 가져온다.

4. 가져온 DP 위치의 path set를 현 path로 update 해준다.

5. DP[0]에서 더 큰값을 출력, 이때의 path를 sort 해준 뒤 출력

 

 

참고로 내 코드는 문제를 풀때 node 번호가 0부터 시작하길 원해서 lambda로 node 값을 수정했는데, 생각해보니 node값을 출력해야돼서 다시 lambda로 node 값을 수정하는 불상사를 겪었다....

 

 

오늘의 교훈) 문제에서 요구하는 사항을 잘 파악하자