Baekjoon Online Judge

algorithm practice

- 단계별 문제풀기

31. 트리에서의 동적 계획법

트리에 동적 계획법을 적용해 봅시다.

Java / Python

4. 우수 마을

1949번

또다른 트리 DP 문제

이번 문제는 각 마을 주민 수와 마을 사이의 길에 대한 정보가 주어졌을 때, 주어진 조건을 만족하도록 '우수 마을'을 선정하는 프로그램을 작성하는 문제이다.

dfs를 이용하며, dp배열은 이차원 배열로 생성하여서 n번 마을이 우수 마을인 경우와 아닌 경우로 나눈다.

• dp[n][1] : n번 마을이 우수 마을일 경우, n번 마을을 루트로 하는 서브트리의 마을 주민 수의 총합
• dp[n][0] : n번 마을이 우수 마을이 아닌 경우, n번 마을을 루트로 하는 서브트리의 마을 주민 수의 총합

• Java

import java.io.*;
import java.util.*;

public class Main {
	static int N;
	static int[] tree;
	static int[][] dp;
	static LinkedList<Integer>[] list;

	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
		StringTokenizer st;

		N = Integer.parseInt(br.readLine());

		dp = new int[N + 1][2];
		list = new LinkedList[N + 1];
		tree = new int[N + 1];

		for (int i = 0; i <= N; i++) {
			list[i] = new LinkedList<>();
		}

		st = new StringTokenizer(br.readLine());
		for (int i = 1; i <= N; i++) {
			tree[i] = Integer.parseInt(st.nextToken());
		}

		for (int i = 0; i < N - 1; i++) {
			st = new StringTokenizer(br.readLine());

			int u = Integer.parseInt(st.nextToken());
			int v = Integer.parseInt(st.nextToken());

			list[u].add(v);
			list[v].add(u);
		}

		dfs(1, -1);

		bw.write(String.valueOf(Math.max(dp[1][0], dp[1][1])) + "\n");
		bw.flush();
		bw.close();
		br.close();
	}

	public static void dfs(int now, int p) {
		for (int next : list[now]) {
			if (next != p) {
				dfs(next, now);
				dp[now][1] += dp[next][0];
				dp[now][0] += Math.max(dp[next][0], dp[next][1]);
			}
		}
		dp[now][1] += tree[now];
	}
}

• Python

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

N = int(input())
person = [0] + list(map(int, input().split()))
tree = [[] for i in range(N + 1)]
visit = [False for i in range(N + 1)]
dp = [[0] * 2 for i in range(N + 1)]

def dfs(start):
    visit[start] = True

    for i in tree[start]:
        if not visit[i]:
            dfs(i)
            dp[start][1] += dp[i][0]
            dp[start][0] += max(dp[i][0], dp[i][1])
    dp[start][1] = dp[start][1] + person[start]
        
for i in range(N - 1):
    u, v = map(int, input().split())
    tree[u].append(v)
    tree[v].append(u)
    
dfs(1)
print(max(dp[1][0], dp[1][1]))