21 Days Training // Day 10

T0. Party

A classic DP on tree.

Notice that :

1. For every leaf node, if it is not selected, the result is dp[u][0] = 0, or else it is dp[u][1]=w[u]
2. For every non-leaf node, if it is not selected, its subs can freely decide whether to be selected, the best result is dp[u][0] = sum(max(dp[v][0], dp[v][1])) (Pseudo Code Warning)
3. For every non-leaf node, if it is selected, its subs cannot be selected, the result is dp[u][1] = w[u] + sum(dp[v][0]) (Pseudo Code Warning).
4. It can be proved that dp[u] does not affect dp[v].

So, we can compute dp for each node, and decide whether to choose dp[root][0] or dp[root][1], the final result is max(dp[root][0], dp[root][1]).

Code :

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N;cin >> N;
    int w[N];
    for (int i=0;i<N;i++) cin >> w[i];
    map<int, set<int> > subs;
    map<int, int> par;
    for (int i=0;i<N-1;i++) {
        int u, v;cin >> u >> v;
        u--;v--;
        if (subs.find(v)==subs.end()) subs[v] = set<int>();
        subs[v].insert(u);
        par[u] = v;
    }
    queue<int> Q;
    int rt;
    for (int i=0;i<N;i++) {
        if (par.find(i)==par.end()) {
            rt=i;break;
        }
    }
    int ind[N];
    for (int i=0;i<N;i++) {
        if (subs.find(i)==subs.end()) {
            ind[i] = 0;
            Q.push(i);
        }
        else {
            ind[i] = subs[i].size();
        }
    }
    int dp[N][2];

    while (!Q.empty()) {
        int cur=Q.front();Q.pop();
        if (rt!=cur) ind[par[cur]]--;
        if (rt!=cur&&ind[par[cur]]==0) Q.push(par[cur]);
        if (subs.find(cur)==subs.end()) dp[cur][0] = 0, dp[cur][1] = w[cur];
        else {
            dp[cur][1] = w[cur];
            dp[cur][0] = 0;
            for (set<int>::iterator it=subs[cur].begin();it!=subs[cur].end();it++) {
                dp[cur][0] += max(dp[*it][0], dp[*it][1]);
                dp[cur][1] += dp[*it][0];
            }
        }
    }
    cout << max(dp[rt][0], dp[rt][1]);
    return 0;
}

Notes :

1. In this problem, we do not know the number of the root, so we need to search for it. (Line 19~24)
2. While calculating dp[u], the order must be from the leaf to the root. (Line 25~35)

Complexity : O(n)