[BOJ] 1949. 우수 마을 (🥇, 트리 DP)

🔗

풀이

트리 형태지만 양방향이기 때문에 루트 노드는 어느 노드가 되어도 상관 없는 구조이다. 문제만 읽으면 마을 하나를 선택하고 그 다음 마을을 건너뛰는 식으로만 하면 될 것 같지만 실제로 예제를 보면 아래와 같은 경우도 가능하다.

2 - 6 은 서로 인접해있지만 둘 다 우수 마을은 아니다. 하지만 각자가 우수 마을과 인접하고 있기 때문에 조건을 만족할 수 있다. 이렇게 네 개의 마을을 선택한 경우가 14000으로 최대가 될 수 있다.

모든 노드는 선택되거나 선택되지 않는 두 가지 경우가 존재할 수 있다. 만약 현재 노드가 선택이 되었다면 현재 노드와 연결된 노드들은 선택되면 안 된다. 반대로 현재 노드가 선택되지 않았다면, 연결된 노드들은 선택될 수도 있고 선택되지 않을 수도 있다. 위의 2-6과 같은 관계처럼 2가 선택되지 않은 경우 6번 노드는 선택될 수도 있고 아닐 수도 있는 것이다. 중요한 것은 최대값을 구하는 것이므로 6을 선택하지 않는 것이 선택하는 것보다 더 높은 값을 가질 수 있다면 선택하지 않는 것이 유리해진다. 예제에서는 7번 노드가 7000이기 때문에 7번 노드를 선택하고 6번을 선택하지 않는 것이 최대값을 가지기에 유리하다.

N개의 행에 대해 2개의 열을 가지는 dp 배열을 선언하고 해당 노드를 선택한 경우 구할 수 있는 인구의 합을 0열에 선택하지 않는 경우를 1열에 저장한다고 가정하자. 탐색의 순서는 트리의 자료구조이기 때문에 루트 노드에서 시작해서 리프 노드까지 탐색한 후 리프 노드에서부터 다시 거꾸로 계산해나오는 방식이 편리하다. 양방향이기 때문에 루트 노드를 어떤 노드로 선정해도 상관이 없고 리프 노드도 딱 잘라 리프 노드라고 하기가 어렵다. 그래서 visited 배열을 통해서 아무 노드나 루트 노드로 잡아서 선택한 후 연결된 노드들로 이동하면서 이미 방문한 노드들에 대해서는 방문 처리를 해서 제외하는 방식으로 리프 노드를 찾는다.

우선 방문한 노드에 대해서는 해당 노드를 선택한 경우는 반드시 해당 노드에 해당하는 인구 수를 합하게 되기 때문에 dp[해당노드][0] += (해당 노드의 인구 수)를 한다. 이후 연결된 노드들(자식 노드)를 돌면서 아래의 점화식을 만족하도록 dp를 채워나간다.

dp[현재 노드][0] += dp[자식 노드][1] (현재 노드를 선택 O, 자식 노드는 선택X)
dp[현재 노드][1] += Math.max(dp[자식 노드][0], dp[자식 노드][1]) (현재 노드를 선택하지 않았다면 자식 노드 중에선 더 큰 값만 선택)

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

public class Main {
    static List<List<Integer>> tree;
    static int[] population;
    static int[][] dp;
    static void dfs(int cur, boolean[] visited) {
        dp[cur][0] += population[cur];
        
        visited[cur] = true;
        for (int child : tree.get(cur)) {
            if (visited[child] == true) continue;
            dfs(child, visited);
            dp[cur][0] += dp[child][1];
            dp[cur][1] += Math.max(dp[child][0], dp[child][1]);
        }
        visited[cur] = false;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        StringTokenizer st = new StringTokenizer(br.readLine());
        population = new int[N];
        for (int i = 0; i < N; i++) {
            population[i] = Integer.parseInt(st.nextToken());
        }

        tree = new ArrayList<>();
        for (int i = 0; i < N; i++) {
            tree.add(new ArrayList<>());
        }

        int a; 
        int b;
        for (int i = 0; i < N-1; i++) {
            st = new StringTokenizer(br.readLine());
            a = Integer.parseInt(st.nextToken());
            b = Integer.parseInt(st.nextToken());
            tree.get(a-1).add(b-1);
            tree.get(b-1).add(a-1);
        }

        dp = new int[N][2];
        boolean[] visited = new boolean[N];
        dfs(0, visited);
        
        System.out.println(Math.max(dp[0][0], dp[0][1]));
    }
}