[Tree/DP] BOJ 2213 트리의 독립집합

문제 링크 : www.acmicpc.net/problem/2213

2213번: 트리의 독립집합

예제의 결과에서 알 수 있듯이, 서브트리의 부모 노드를 선택하지 않는다고 해서 반드시 자식 노드가 선택되는 것은 아니다. 이는 곧 그리디 알고리즘으로는 해결할 수 없다는 것을 의미한다.

따라서 Tree 구조에서의 DP 문제로 접근해 문제를 풀어보자.

어떤 한 서브트리에서, 문제를 두 가지 경우로 나눌 수 있다.

 

① 서브트리의 root 노드를 선택했을 경우 (코드의 sel = 1)

이 경우, 모든 자식 노드들은 선택할 수 없다.

따라서 인접한 모든 자식 노드들을 선택하지 않았을 때의 최대 가중치들을 더해준다.

이 때, 현재 root 노드를 선택했으므로 노드의 가중치를 추가로 더해준다.

 

② 서브트리의 root 노드를 선택하지 않았을 경우 (sel = 0)

이 경우, 각 자식 노드들은 선택될 수도 있고 선택되지 않을 수도 있다.

인접한 자식을 선택했을 때, 그 자식 노드의 자식 노드 가중치가 더 높을 수 있기 때문이다.

따라서, 인접 자식 노드를 선택한 경우와 선택하지 않은 경우를 모두 고려해 가중치가 더 큰 값을 선택한다.

 

위 두 계산을 각각의 부분 트리에 대해 재귀적으로 실행해 원하는 답을 찾는다.

이 때의 최종 결과는 root을 선택한 경우와 그렇지 않은 경우로 나뉘므로, 결과값이 더 큰 경우에 대해

다시 재귀적으로 살펴보며 집합에 포함되는 노드 번호를 추가하면 된다.

 

[코드]

#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <vector>
#include <algorithm>
 
using namespace std;
typedef pair<int, int> p;
 
int n, val[10001];
vector<int> adj[10001], result;
int dp[10001][2];
 
int get_dp(int node, int pnode, int sel) {
    if (dp[node][sel]) return dp[node][sel];        //Memoization
 
    //sel이 1이면 현재 노드를 선택한다는 의미
    int ret = 0;
    if (sel) {
        //리프 노드의 경우
        if (node != 1 && adj[node].size() == 1) return dp[node][sel] = val[node];
 
        //현재 노드를 선택한다면, 인접 자식 노드를 선택할 수 없다.
        for (int i = 0; i < adj[node].size(); i++) {
            int child = adj[node][i];
            if (child == pnode) continue;
 
            ret += get_dp(child, node, 0);
        }
 
        //현재 노드를 선택했으므로 가중치를 더해 반환
        ret += val[node];
    }
 
    else {
        //리프 노드의 경우
        if (node != 1 && adj[node].size() == 1) return 0;
 
        //현재 노드를 선택하지 않는다면, 모든 인접 노드를 선택'할 수' 있다.
        for (int i = 0; i < adj[node].size(); i++) {
            int child = adj[node][i];
            if (child == pnode) continue;
 
            //인접 노드를 선택한 경우와 그렇지 않은 경우 중 더 큰 값을 더함
            //선택을 했어도 더할 수 있고, 그렇지 않아도 더할 수 있음
            int child_sel = get_dp(child, node, 1);
            int child_notsel = get_dp(child, node, 0);
            ret += max(child_sel, child_notsel);
        }
    }
 
    return dp[node][sel] = ret;
}
 
void get_group(int node, int pnode, int sel) {
    if (sel) {
        //현재 노드를 선택하므로, 최종 배열에 삽입
        result.push_back(node);
 
        for (int i = 0; i < adj[node].size(); i++) {
            int child = adj[node][i];
            if (child == pnode) continue;
 
            get_group(child, node, 0);
        }
    }
 
    else {
        for (int i = 0; i < adj[node].size(); i++) {
            int child = adj[node][i];
            if (child == pnode) continue;
 
            //현재 노드를 선택하지 않으므로, 두 경우 모두 확인
            int child_sel = dp[child][1];
            int child_notsel = dp[child][0];
 
            if (child_sel > child_notsel) get_group(child, node, 1);
            else get_group(child, node, 0);
        }
    }
}
 
void init() {
    int f, s;
    cin >> n;
 
    for (int i = 1; i <= n; i++) cin >> val[i];
 
    for (int i = 0; i < n - 1; i++) {
        cin >> f >> s;
        adj[f].push_back(s);
        adj[s].push_back(f);
    }
}
 
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(NULL);
 
    init();
    int root_sel = get_dp(1, 0, 1);
    int root_notsel = get_dp(1, 0, 0);
 
    printf("%d\n", max(root_sel, root_notsel));
 
    if (root_sel > root_notsel) get_group(1, 0, 1);
    else get_group(1, 0, 0);
    
    sort(result.begin(), result.end());
    for (int i = 0; i < result.size(); i++) printf("%d ", result[i]);
    printf("\n");
 
    return 0;
}