[트리 DP(Top-Down)] 백준 2213 트리의 독립집합 C++

 

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

2213번: 트리의 독립집합

 

#include<iostream>
#include<vector>
#include<algorithm>
#define MAX 10001
using namespace std;

int n, a, b, ans;
int dp[MAX][2];
int visited[MAX];
int check[MAX];
vector<int> w;
vector<int> v[MAX];
vector<int> tree[MAX];
vector<int> r;

void dfs(int x) {
    visited[x] = 1;
    
    for (int i = 0; i < v[x].size(); i++) {
        int temp = v[x][i];
        
        if (visited[temp] == 1) continue; // 조심
        tree[x].push_back(temp);
        
        dfs(temp);
    }
}

void calDp(int x) {
    for (int i = 0; i < tree[x].size(); i++) {
        int temp = tree[x][i];

        calDp(temp);
        dp[x][0] += dp[temp][1];
        dp[x][1] += max(dp[temp][0], dp[temp][1]);   
    }
    dp[x][0] += w[x];
}

void tracking(int s, int o) {
    if (o == 0) {
        r.push_back(s);
        for (int i = 0; i < tree[s].size(); i++) {
            int temp = tree[s][i];
            tracking(temp, 1);
        }
    }
    else {
        for (int i = 0; i < tree[s].size(); i++) {
            int temp = tree[s][i];
            if (dp[temp][1] >= dp[temp][0]) tracking(temp, 1);
            else tracking(temp, 0);
        }
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    cin >> n;
    
    for (int i = 0; i < n; i++) {
        cin >> a;
        w.push_back(a);
    }

    for (int i = 0; i < n - 1; i++) {
        cin >> a >> b;
        a--;
        b--;
        v[a].push_back(b);
        v[b].push_back(a);
    }
    
    dfs(0);
    calDp(0);

    if (dp[0][0] >= dp[0][1]) {
        ans = dp[0][0];
        tracking(0, 0);//포함
    }
    else {
        ans = dp[0][1];
        tracking(0, 1);
    }
    
    sort(r.begin(), r.end());
    cout << ans << "\n";
    for (int i = 0; i < r.size(); i++) cout << r[i]+1 << " ";
}