BOJ 2213 트리의 독립집합

https://www.acmicpc.net/problem/2213

문제

그래프 G(V, E)에서 정점의 부분 집합 S에 속한 모든 정점쌍이 서로 인접하지 않으면 (정점쌍을 잇는 간선이 없으면) S를 독립 집합(independent set)이라고 한다. 독립 집합의 크기는 정점에 가중치가 주어져 있지 않을 경우는 독립 집합에 속한 정점의 수를 말하고, 정점에 가중치가 주어져 있으면 독립 집합에 속한 정점의 가중치의 합으로 정의한다. 독립 집합이 공집합일 때 그 크기는 0이라고 하자. 크기가 최대인 독립 집합을 최대 독립 집합이라고 한다.

문제는 일반적인 그래프가 아니라 트리(연결되어 있고 사이클이 없는 그래프)와 각 정점의 가중치가 양의 정수로 주어져 있을 때, 최대 독립 집합을 구하는 것이다.

---

입력

첫째 줄에 트리의 정점의 수 n이 주어진다. n은 10,000이하인 양의 정수이다. 1부터 n사이의 정수가 트리의 정점이라고 가정한다. 둘째 줄에는 n개의 정수 $w_1, w_2, …, w_n$이 주어지는데, $w_i$는 정점 i의 가중치이다(1 ≤ i ≤ n). 셋째 줄부터 마지막 줄까지는 간선의 리스트가 주어지는데, 한 줄에 하나의 간선을 나타낸다. 간선은 정점의 쌍으로 주어진다. 입력되는 정수들 사이에는 빈 칸이 하나 있다. 가중치들의 값은 10,000을 넘지 않는 자연수이다.

---

출력

첫째 줄에 최대 독립집합의 크기를 출력한다. 둘째 줄에는 최대 독립집합에 속하는 정점을 오름차순으로 출력한다. 최대 독립 집합이 하나 이상일 경우에는 하나만 출력하면 된다.

---

예제 입력 1

7
10 30 40 10 20 20 70
1 2
2 3
4 3
4 5
6 2
6 7

예제 출력 1

140
1 3 5 7

---

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;

public class Main {
    static int N;
    static int[] W;

    // 최대 독립 집합으로 선택된 정점을 찾기 위한 ArrayList
    static ArrayList<Integer> roots = new ArrayList<>();

    // 트리의 간선 정보를 저장하기 위한 ArrayList[]
    static ArrayList<Integer>[] tree;
    static int[][] dp;

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

        // 간선 정보를 입력받을 때, 두 정점 중 작은 정점을 기준으로 간선 정보를 저장함
        for (int i = 1; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken()), b = Integer.parseInt(st.nextToken());
            if(a < b) tree[a].add(b);
            else tree[b].add(a);
        }
        // dp의 첫 번째 인덱스는 정점의 번호를 의미, 두 번째 인덱스는 해당 정점의 선택 여부를 의미 (0은 선택X, 1은 선택O)
        dp = new int[N + 1][2];
        for (int i = 1; i <= N; i++) {
            /**
             * i번 째 정점에 저장된 간선 정보가 없다는 것은 dp로 검사할 때 마지막 정점이기 때문에
             * dp[i][0]은 해당 정점이 선택되지 않았기 때문에 0으로 두고
             * dp[i][1]은 해당 정점의 가중치로 초기화함
             * i번 째 정점에 저장된 간선 정보가 있으면 아직 검사하지 않았다는 의미로 -1로 초기화
             */
            if(tree[i].size() == 0) dp[i][1] = W[i];
            else dp[i][0] = dp[i][1] = -1;
        }

        // 첫 번째 정점을 선택하지 않았을 때와 선택했을 때의 가중치 중 높은 것이 정답
        int max = Integer.max(getMaxW(1, 0), getMaxW(1, 1));
        // 최대 독립집합에 속하는 정점들 서칭
        findRoot(1, dp[1][0] > dp[1][1] ? 0 : 1);
        Collections.sort(roots);

        StringBuilder answer = new StringBuilder();
        answer.append(max).append('\n');
        for (int root: roots) answer.append(root).append(' ');
        System.out.println(answer);
    }

    static int getMaxW(int idx, int isSelected) {
        if(dp[idx][isSelected] != -1) return dp[idx][isSelected];

        dp[idx][isSelected] = 0;
        // 해당 정점에 간선으로 연결된 정점들을 검사
        for (int i = 0; i < tree[idx].size(); i++) {
            int checkIdx = tree[idx].get(i);
            // 현재 정점이 선택되지 않았으면 연결된 정점들이 선택됐을 때와 선택되지 않았을 때의 가중치 모두 검사
            if(isSelected == 0)
                dp[idx][isSelected] += Integer.max(getMaxW(checkIdx, 0), getMaxW(checkIdx, 1));
            // 현재 정점이 선택되었으면 연결된 정점들이 선택되지 않았을 때의 가중치만 검사
            else
                dp[idx][isSelected] += getMaxW(checkIdx, 0);
        }
        if(isSelected == 1) dp[idx][isSelected] += W[idx]; // 현재 정점이 선택되었으면, 현재 정점의 가중치+

        return dp[idx][isSelected];
    }

    /**
     * 기본적으로 정점들이 선택되지 않았을 때와 선택되었을 때의 최대 가중치를 비교해
     * 높은 가중치가 있는 루트로 검사하는 방식이다.
     */
    static void findRoot(int idx, int isSelected) {
        if(isSelected == 1) {
            roots.add(idx);
            for (int i = 0; i < tree[idx].size(); i++) findRoot(tree[idx].get(i), 0);
            return;
        }

        for (int i = 0; i < tree[idx].size(); i++) {
            int checkIdx = tree[idx].get(i);
            findRoot(checkIdx, dp[checkIdx][0] > dp[checkIdx][1] ? 0 : 1);
        }
    }
}