1949. 우수 마을 (Python)

1949번: 우수 마을

 

풀이

 임의의 마을 하나를 루트 노드로 설정하여 해당 노드부터 리프 노드까지 탐색하면 된다. 코드에서 사용된 dp[n][flag]의 경우 flag가 0일 경우는 n번 마을을 우수 마을로 설정하지 않았을 때의 마을 주민 수의 최댓값이고 flag가 1일 경우는 n번 마을을 우수 마을로 설정했을 때의 마을 주민 수의 최댓값을 의미한다. 

 

 재귀 호출되는 함수의 경우 3개의 매개변수를 받게 되어있는데 node의 경우는 현재 마을 번호, parent의 경우는 현재 마을의 부모 마을 번호, parentCheck의 경우는 부모 마을을 우수 마을로 선정했는지 여부를 나타낸다.

def solve(node, parent, parentCheck)

 

 만약 부모 노드를 우수 마을로 선정했다면 현재 마을은 우수 마을로 선정할 수 없기 때문에 다음 마을에 대하여 함수를 호출할 때 parentCheck를 0으로 설정하면 되고, 부모 노드를 우수 마을로 선정하지 않았다면 현재 마을을 우수 마을로 선정할 수 있기 때문에 parentCheck를 1로 설정하고 함수를 호출하면 된다.

 

코드

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

# node: 현재 마을 번호, parent: 부모 마을 번호, parentCheck: 부모 마을 우수 마을 선정 여부 
# parentCheck -> 0: 부모 노드를 우수 마을로 선정하지 않음, 1: 부모 노드를 우수 마을로 선정함
def solve(node, parent, parentCheck):
    # [현재 마을][부모 노드 우수 마을 여부]에 대하여 이미 값이 들어가 있다면 바로 함수 종료 
    if dp[node][parentCheck] != -1: 
        return dp[node][parentCheck]
    
    dp[node][parentCheck] = 0
    notUse = 0 # 현재 마을을 우수 마을로 선정하지 않았을 때
    use = 0    # 현재 마을을 우수 마을로 선정했을 때

    if parentCheck == 0: # 부모 마을이 우수 마을이 아니라면
        use = people[node] # 현재 마을의 주민 수를 더해놓는다.

    for p in path[node]: # 현재 마을과 이어진 마을에 대하여 탐색
        if p == parent: # 부모 마을의 경우는 건너뛴다.
            continue 
        notUse += solve(p, node, 0) # 현재 마을을 우수 마을로 설정하지 않고 함수를 호출한다.
        
        if parentCheck == 0: # 부모 마을이 우수 마을이 아니라면
            use += solve(p, node, 1) # 현재 마을을 우수 마을로 설정하고 함수를 호출한다.

    dp[node][parentCheck] = max(notUse, use) # 두 값중 큰 값을 배열에 넣는다.
    return dp[node][parentCheck]



if __name__ == '__main__':
    N = int(input()) # 마을의 개수
    people = [0] + list(map(int, input().split())) # 마을 주민 수
    path = [[] for _ in range(N+1)]
    dp = [[-1, -1] for _ in range(N+1)]

    for _ in range(N-1):
        a, b= map(int, input().split())
        path[a].append(b)
        path[b].append(a)
    
    res = solve(1, 0, 0)
    print(res)