Tree with Maximum Cost CodeForces - 1092F （换根DP）

题目

https://cn.vjudge.net/problem/CodeForces-1092F#author=0

题意

给你个树 让你选定一个根 使得所有点权值*到根距离 的和最大

思路

换根DP 思考换根后改变了什么

比如存在边 u->v  当前u为根

换根后 v的子树所有点减少了一次 除v子树所有点增加了一次

那就预处理sum[i] 表示i子树所有点的权值

换根

先减去换跟前sum[v] 

换根后 sum[u]变为了 sum[u] - sum[v]            sum[v]变为了sum[v] + sum[u]

再加上换根后的sum[u] 即换根后的权值和

代码

#include <bits/stdc++.h>

using namespace std;
const int maxn = 2e5 + 100;
typedef long long ll;
int a[maxn];
vector<int> G[maxn];
ll val;
ll sum[maxn];
ll ans;
void dfs(int u,int d,int fa)
{
    sum[u] = (ll)a[u];
    val += 1ll*d*a[u];
    int len = G[u].size();
    for(int i = 0;i < len;i++)
    {
        int v = G[u][i];
        if(v == fa) continue;
        dfs(v,d+1,u);
        sum[u] += sum[v];
    }
}
void dfs2(int u,int fa)
{
    ans = max(ans,val);
    int len = G[u].size();
    for(int i = 0;i < len;i++)
    {
        int v = G[u][i];
        if(v == fa) continue;

        val -= sum[v];
        sum[u] -= sum[v];
        sum[v] += sum[u];
        val += sum[u];

        dfs2(v,u);

        val -= sum[u];
        sum[v] -= sum[u];
        sum[u] += sum[v];
        val += sum[v];

    }
}
int main()
{
    ios::sync_with_stdio(false);
    int n;
    cin>>n;
    for(int i = 1;i <= n;i++)
    {
        cin>>a[i];
        G[i].clear();
    }
    for(int i = 1;i < n;i++)
    {
        int u,v;
        cin>>u>>v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    val = 0;
    dfs(1,0,-1);
    ans = val;
    dfs2(1,-1);

    cout<<ans<<endl;
    return 0;
}