Codeforces Round #384 (Div. 2)D. Chloe and pleasant prizes（树DP）

Posted 2022-12-12 AC_Arthur

题目链接：点击打开链接

思路：

比较简单的树DP， 用dp[u][id]表示当前以u为根的子树还已经找到几个子树的最大值。  转移比较多， 一方面可以转移到某一个儿子， 表示问题在以后解决， 一方面如果id==1说明还要找1个子树，可以直接用val[u]更新， val[u]表示该子树的和。   如果id == 0说明还要找两个子树， 我们用两个最大的儿子值更新即可。

细节参见代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>
#include <stack>
#include <ctime>
#include <bitset>
#include <cstdlib>
#include <cmath>
#include <set>
#include <list>
#include <deque>
#include <map>
#include <queue>
#define Max(a,b) ((a)>(b)?(a):(b))
#define Min(a,b) ((a)<(b)?(a):(b))
using namespace std;
typedef long long ll;
typedef long double ld;
const double eps = 1e-6;
const double PI = acos(-1);
const int mod = 1000000000 + 7;
const int seed = 131;
const ll INF = ll(1e16);
const int maxn = 2e5 + 10;
int n, vis[maxn][3], kase = 0, par[maxn];
ll a[maxn], d[maxn][3], val[maxn];
vector<int> g[maxn];
ll dfs(int u, int fa) 
    int len = g[u].size();
    val[u] = a[u];
    par[u] = fa;
    for(int i = 0; i < len; i++) 
        int v = g[u][i];
        if(v == fa) continue;
        val[u] += dfs(v, u);
    
    return val[u];

ll dp(int u, int id) 
    ll& ans = d[u][id];
    if(vis[u][id] == kase) return ans;
    vis[u][id] = kase;
    ans = -INF;
    int len = g[u].size();
    vector<ll> cur;
    for(int i = 0; i < len; i++) 
        int v = g[u][i];
        if(v == par[u]) continue;
        ll di = dp(v, id);
        if(di > -INF) ans = max(ans, di);
        if(id == 0) cur.push_back(dp(v, id+1));
    
    sort(cur.begin(), cur.end());
    len = cur.size();
    if(id == 1) ans = max(ans, val[u]);
    if(id == 0 && len > 1) ans = max(ans, cur[len-1]+cur[len-2]);
    return ans;

int main() 
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) scanf("%I64d", &a[i]);
    for(int i = 1; i < n; i++) 
        int u, v; scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    
    ++kase;
    dfs(1, 1);
    ll ans = dp(1, 0);
    if(ans == -INF) printf("Impossible\\n");
    else printf("%I64d\\n", ans);
    return 0;