[파이썬] 백준 1949 : 우수마을 (골드2)

[파이썬] 백준 1949 : 우수마을 (골드2)

1949번: 우수 마을

---

문제

---

풀이

0. 방향성 생각

DFS를 사용해서 루트노드에서 시작해서 리프노드부터 DP 시작을 해준다.

DFS를 마치고 나면 DP 루트노드를 확인하면 가중치 합이 저장되게 작성한다.

상태는 두 가지 상태가 존재한다. dp[state][노드수]로 dp 테이블 작성하기

1. 입력

import sys
sys.setrecursionlimit(10**6)

input = lambda : sys.stdin.readline().rstrip()
T = int(input())
population = dict(zip(range(1,T+1),list(map(int,input().split()))))

tree = [[] for _ in range(T+1)]
for _ in range(T-1):
    a,b = map(int,input().split())
    tree[b].append(a)
    tree[a].append(b)

딕셔너리 zip으로 한 번에 받아준다.

tree에 양방향 간선 추가.

 

2. DFS 함수 정의, 출력

visit = set()
dp = [[0]*(T+1) for _ in range(2)]
dp[1] = [0] + list(population.values())

def dfs(x):
    global visit,tree
    visit.add(x)
    for nx in tree[x]:
        if nx not in visit:
            dfs(nx)
            dp[0][x] += max(dp[0][nx],dp[1][nx])
            dp[1][x] += dp[0][nx]

dfs(1)
print(max(dp[i][1] for i in range(2)))

DP[선정유무][노드]로 dp 테이블을 정의한다. 0번은 더미노드라서 T+1개로 선정

 

DFS를 통해 계속 내려가면 리프노드에 도달한다.

 

리프노드에 도달한 후 자식노드 nx를 어떤 상태로 지정하냐에 따라서 부모노드 x의 값이 변한다

dp[0][x] (부모노드를 우수마을로 선정하지 않은 경우) : 자식노드 선정, 선정하지 않는 경우 중 큰 것으로

o x x o 같은 경우가 생긴다. x가 연속으로 등장 가능

dp[1][x] (부모노드를 우수마을로 선정하는 경우) : 자식노드는 선정하지 않는다.

 

---

전체코드

import sys
sys.setrecursionlimit(10**6)

input = lambda : sys.stdin.readline().rstrip()
T = int(input())
population = dict(zip(range(1,T+1),list(map(int,input().split()))))

tree = [[] for _ in range(T+1)]
for _ in range(T-1):
    a,b = map(int,input().split())
    tree[b].append(a)
    tree[a].append(b)

visit = set()
dp = [[0]*(T+1) for _ in range(2)]
dp[1] = [0] + list(population.values())

def dfs(x):
    global visit,tree
    visit.add(x)
    for nx in tree[x]:
        if nx not in visit:
            dfs(nx)
            dp[0][x] += max(dp[0][nx],dp[1][nx])
            dp[1][x] += dp[0][nx]

dfs(1)
print(max(dp[i][1] for i in range(2)))

코멘트

분명 맞는데 왜 안되나 했는데 우수마을 선정 안한쪽을 활성화시켜놔서,,, 후.....................