每日亿题--2021.5.11换根dp

换根dp

之前发过换根dp通用做法的博文，不记得了可以去看看树形dp–换根

今天这题是一道常见的换根dp的题目，是换根dp的简单拓展，两个dfs即可解决。
注意数据范围，1≤n,ai≤2×10^5，考虑到转移方程计算的时候存在乘法，为了防止爆int，所以我全程采用了longlong。感觉也有点前缀和的思想。

思路：
1、知道每条边的权值为1，且每次以某一结点为根节点时，我们求的都是其他节点到根节点的边数*该节点的权值，所以需要用dep数组记录一下每个节点间的深度关系。（深度关系是我自己命名的叫法，具体意思可以去理解代码，其实就是深度）
2、转移方程。我们会发现当根节点为u时，向节点v转移 ，即使节点v为根节点，这时，以v为根节点的子树中，所有的节点相对于初始根节点u而言，都离现在根节点v近了一条边，所以要减去以v为根节点的这一部分子树的权值和的一倍。同理，其他节点的权值和都的加上一倍。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll N=2e5+5;
// son[i]表示以节点i为根节点的子树的权值和
// dep[i]表示i节点的深度，这里我将树看成以节点1为根节点的一棵树，所以dep[1]为1
// dp[i]表示以i为目标节点的值
ll n,a[N],ans;
ll to[N*2],nex[N*2],fir[N],idx;
ll son[N],dep[N],dp[N];

void add(int u,int v){
    to[++idx]=v;
    nex[idx]=fir[u];
    fir[u]=idx;
}

void dfs1(int u,int fa){
    son[u]=a[u],dep[u]=dep[fa]+1;
    for(int i=fir[u];i;i=nex[i]){
        int v=to[i];
        if(v!=fa){
            dfs1(v,u);
            son[u]+=son[v];
        }
    }
}

void dfs2(int u,int fa){
    for(int i=fir[u];i;i=nex[i]){
        int v=to[i];
        if(v!=fa){
            dp[v]=dp[u]+(son[1]-son[v])-son[v];
            dfs2(v,u);
        }
    }
}

int main(){
    scanf("%lld",&n);
    for(int i=1;i<=n;i++) scanf("%lld",&a[i]);
    for(int i=1;i<n;i++){
        int u,v; scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    dfs1(1,0);
    for(int i=2;i<=n;i++) dp[1]=dp[1]+(dep[i]-dep[1])*a[i];
    dfs2(1,0);
    for(int i=1;i<=n;i++) ans=max(ans,dp[i]);
    cout<<ans<<"\n";
    return 0;
}