CF-1092 F. Tree with Maximum Cost(换根DP)

题目链接

题意：给定一棵树， 每个点有一个点权 a[i] ， 边权为 1 ，表达式

其中，v 是树上的一个节点， dist(x, y)表示点 x 和 y 之间的距离，求此表达式的最大值

代码

#include <bits/stdc++.h>
#define debug(x) cout<<#x<<"="<<x<<endl;
#define debug2(x, y) cout<<#x<<"="<<x<<" "<<#y<<"="<<y<<endl;
#define debug3(x, y, z) cout<<#x<<"="<<x<<" "<<#y<<"="<<y<<" "<<#z<<"="<<z<<endl;
const int maxn = 2e5+50;
using namespace std;
typedef long long ll;
int n;
ll a[maxn], dp[maxn], he[maxn];
vector <int> maps[maxn];

void dfs(int u, int fa){
    he[u] = 0;
    for(int i=0, len=maps[u].size(); i<len; i++){
        int v = maps[u][i];
        if(v != fa){
            dfs(v, u);
            dp[u] += dp[v] + he[v];
            he[u] += he[v];
        }
    }
    he[u] += a[u];
}
ll ans;
void dfs1(int u, int fa){
    for(int i=0, len=maps[u].size(); i<len; i++){
        int v = maps[u][i];
        if(v != fa){
            dp[v] =  dp[u] - he[v] + he[1] - he[v];
            ans = max(ans, dp[v]);
            dfs1(v, u);
        }
    }
}

int main()
{
    scanf("%d", &n);
    for(int i=1; i<=n; i++)
        scanf("%d", &a[i]);
    for(int i=1, a, b; i<n; i++){
        scanf("%d%d", &a, &b);
        maps[a].push_back(b);
        maps[b].push_back(a);
    }
    dfs(1, 1);
    ans = dp[1];
    dfs1(1, 1);
    printf("%lld
", ans);
    return 0;
}