树上差分,换根dp——2020-camp-day1-E

#include<bits/stdc++.h>
#include<vector>
using namespace std;
#define N 300005
#define ll long long 
vector<int>G[N];

ll n,m;
pair<int,int>p[N];

ll fa[N],size[N],deep[N],son[N],top[N],ind,id[N];
void dfs1(int x,int pre,int dep){
    size[x]=1;deep[x]=dep;
    for(auto y:G[x]){
        if(y==pre)continue;
        fa[y]=x;
        dfs1(y,x,dep+1);
        size[x]+=size[y];
        if(size[son[x]]<size[y])son[x]=y;
    }
}
void dfs2(int x,int tp){
    top[x]=tp;id[x]=++ind;
    if(son[x])dfs2(son[x],tp);
    for(auto y:G[x]){
        if(y!=son[x] && y!=fa[x])dfs2(y,y);
    }
}
int lca(int x,int y){
    while(top[x]!=top[y]){
        if(deep[top[x]]<deep[top[y]])swap(x,y);
        x=fa[top[x]];
    }
    //在同一条重链上了 
    if(deep[x]>deep[y])swap(x,y);
    return x;
}

ll A[N],B[N],ans[N];
void dfs3(int u,int pre){
    for(auto v:G[u]){
        if(v==pre)continue;
        dfs3(v,u);
        A[u]+=A[v];B[u]+=B[v];
    }
}
void dfs4(int u,int pre){
    ans[u]+=deep[u]*A[u]+B[u];
    for(auto v:G[u]){
        if(v!=pre)
            ans[v]=ans[u],dfs4(v,u);
    }
}

int main(){
    cin>>n>>m;
    for(int i=1;i<n;i++){
        int u,v;cin>>u>>v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    for(int i=1;i<=m;i++)
        cin>>p[i].first>>p[i].second;
    
    dfs1(1,0,1);
    dfs2(1,1);

    /*
    处理每条路径(u,v),路径的顶点为fa,路径长度是len 
    这条路径对u->fa上每条边的贡献是 2*(deep[u]-deep[i])+1-len=2*deep[u]+1-len-2*deep[i] 
    因为是等差序列不好直接在树上差分,
    拆开括号,观察式子可以发现除了常数项1-len+2*deep[u]外,另一项之和点的深度有关,而深度前面的那个改变的系数也是个常数
    每条边的总用 A*deep[i]+B 来表示 
    更新差分时先把(u,fa)这一段的B加上 1-len+2*deep[u],然后把这一段的A系数加上-2  
    */ 
    for(int i=1;i<=m;i++){
        int u=p[i].first,v=p[i].second,fa=lca(u,v);
        ll len=deep[u]+deep[v]-deep[fa]*2;
        B[u]+=1-len+2*deep[u];B[fa]-=1-len+2*deep[u];
        B[v]+=1-len+2*deep[v];B[fa]-=1-len+2*deep[v];
        A[u]-=2;A[fa]+=2;A[v]-=2;A[fa]+=2;
        
        ans[1]+=(deep[u]-deep[fa])*(len-deep[u]+deep[fa]);
    }
    dfs3(1,0);
    dfs4(1,0);
    for(int i=1;i<=n;i++)cout<<ans[i]<<'\n';
} 

posted on 2020-01-15 16:31  zsben  阅读(126)  评论(0)    收藏  举报

导航