트리의 독립집합

백준 2213번: 트리의 독립집합

아이디어

루트부터 내려가면서 해당 노드를 선택할지 말지 결정한다. 만약 부모 노드가 선택된 상태라면 현재 노드를 선택하지 않는 경우만 고려하면 되고, 부모 노드가 선택되지 않은 상태라면 현재 노드를 선택하는 경우, 선택하지 않는 경우 두 가지 모두 고려해야 한다.

cache[현재 노드][부모 노드 선택 여부]에 해당 노드로 만들 수 있는 서브트리에서 만들어지는 독립집합 가중치의 최대를 메모이제이션 한다.

독립집합에 속하는 노드을 알아내는 방법은 부모 노드가 선택되지 않았으면서 cache[현재 노드][0] > cache[현재 노드][1]인 경우 선택된 노드이니 vector에 추가해놓고 나중에 sort하여 출력한다.

코드

#include <bits/stdc++.h>

using namespace std;

constexpr int MAX = 10001;
int N;
int w[MAX];
vector<int> adj[MAX], tree[MAX];
bool visited[MAX];
int cache[MAX][2];
vector<int> ans;

int solve(int cur, int p) { // 현재 노드, 부모 노드 선택 여부
    int& ret = cache[cur][p];

    if (!tree[cur].size()) {
        if (p) return ret = 0;
        else return ret = w[cur];
    }

    if (ret != -1) return ret;

    int ret1 = 0, ret2 = 0;
    for (int x : tree[cur]) {
        ret1 += solve(x, 0);
    }
    if (!p) {
        ret2 += w[cur];
        for (int x : tree[cur]) {
            ret2 += solve(x, 1);
        }
    }
    ret = max(ret1, ret2);
    return ret;
}

void track(int cur, int p) {
    if (!p && cache[cur][0] > cache[cur][1]) {
        ans.push_back(cur);
        for (int x : tree[cur]) {
            track(x, 1);
        }
    }
    else {
        for (int x : tree[cur]) {
            track(x, 0);
        }
    }
}

void make_tree(int cur) {
    visited[cur] = true;
    for (int x : adj[cur]) {
        if (!visited[x]) {
            visited[x] = true;
            tree[cur].push_back(x);
            make_tree(x);
        }
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> N;
    for (int i = 1; i < N+1; i++) {
        cin >> w[i];
    }
    for (int i = 0; i < N-1; i++) {
        int a, b;
        cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    
    adj[0].push_back(1);
    make_tree(0);

    memset(cache, -1, sizeof(cache));
    cout << solve(0, 0) << '\n';

    track(1, 0);
    sort(ans.begin(), ans.end());
    for (int x : ans) {
        cout << x << ' ';
    }

    return 0;
}

여담

트리를 만들 때 노드 번호 순서대로 인접 리스트를 참조하여 트리를 만들면 꼬일 수 있다. dfs로 만들자.
또 1번 노드(루트로 설정함)가 포함되는 경우, 포함되지 않는 경우 모두 tracking할 수 있도록 0번 노드를 만들어서 1번 노드와 연결하고 0번 노드를 루트로 해서 계산했다.