CF #527 (Div. 3) F. Tree with Maximum Cost

F. Tree with Maximum Cost

题目链接： Tree with Maximum Cost

题意

给你一棵N个点的树，每个点有一个权值A[i]，现在定义一个MAX值为选定一个点v
$$\sum _{i=1}^{n}dist(i,v)\ast A[i]$$

问这个值最大为多少？

数据范围： $N<=2\ast 10^5;A[i]<=2\ast 10^5$

---

思路

假定这个点是一，那么我们可以把这颗树变为以一为根的树，现在每个点对一的贡献为他的深度*他的值，那么我们把这个分散到他们的中间节点，那么每个点所代表的值就是他子树的之和。
$$ans=\sum _{i=1}^{n}dist(i,v)\ast A[i]=\sum _{i=1}^n点i的所有子节点和$$
所以我们可以先进行一次dfs来找到每个点的子节点之和，从点一开始，进行转移，每次转移只需要 $O(1)$的时间

转移公式 now + (sum[1]-sum[v]) - sum[v];

---

代码

#include <bits/stdc++.h>
using namespace std;
#define rep(i,j,k) for(int i = (int)j;i <= (int)k;i ++)
#define debug(x) cerr<<#x<<":"<<x<<endl
#define pb push_back

typedef long long ll;
const int MAXN = (int)1e6+7;

int N;
ll A[MAXN];
ll sum[MAXN];
ll ans;
vector<int> G[MAXN];

void dfs(int u,int fa) {
    sum[u] = A[u];
    rep(i,0,G[u].size()-1) {
        int to = G[u][i];
        if (to == fa) continue;
        dfs(to,u);
        sum[u] += sum[to];
    }
}

void getMax(int u,int fa,ll now) {
    ans = max(ans,now);
    rep(i,0,G[u].size()-1) {
        int v = G[u][i];
        if (v == fa) continue;
        ll tmp = now + (sum[1]-sum[v]) - sum[v];
        getMax(to,u,tmp);
    }
}

int main()
{
    ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

    cin >> N;
    rep(i,1,N) cin >> A[i];
    rep(i,1,N-1) {
        int u,v;
        cin >> u >> v;
        G[u].pb(v);
        G[v].pb(u);
    }

    dfs(1,-1);
    rep(i,2,N) ans += sum[i];
    getMax(1,-1,ans);
    cout << ans << endl;
}

/*
现在给了我一个树
1.预处理出一个树，
每个点都储存着包括自己子树的值和，那么每次转移只需要减去父节点树的值，加上子节点树的值

*/