牛客练习赛55 -树-所有点对距离的平方和

题目链接：https://ac.nowcoder.com/acm/contest/2927/E

思路：

对于节点4来说，它到其他点的距离dis[i] $$\sum dis(4,i{)}^2=dis[1{]}^2+dis[2{]}^2+dis[3{]}^2+dis[4{]}^2+dis[5{]}^2+dis[6{]}^2+dis[7{]}^2+dis[8{]}^2$$
如果根据dfs转移给子树的时候能够开始求出这个值，那么就可以dfs一次得到答案。
对于下一个点5
$$\sum dis(5,i{)}^2=(dis[1]+1{)}^2+(dis[2]+1{)}^2+(dis[3]+1)+(dis[4]+1)+(dis[5]-1{)}^2+(dis[6]-1{)}^2+(dis[7]-1{)}^2+(dis[8]-1{)}^2$$
$$\sum dis(5,i{)}^2=dis[1{]}^2+dis[2{]}^2+dis[3{]}^2+dis[4{]}^2+dis[5{]}^2+dis[6{]}^2+dis[7{]}^2+dis[8{]}^2$$
$$+2(dis[1]+dis[2]+dis[3]+dis[4])-2(dis[5]+dis[6]+dis[7]+dis[8])+n$$

$$\sum dis(5,i{)}^2=\sum dis(4,i{)}^2+2(dis[1]+dis[2]+dis[3]+dis[4])-2(dis[5]+dis[6]+dis[7]+dis[8])+n$$
我们发现：

$$s1=dis[5]+dis[6]+dis[7]+dis[8]是5的子树(包含5)到4的距离=5所有子树到5的距离和+4的子节点个数:一次dfs就可以求所有节点的s1$$

$$s2=dis[5]+dis[6]+dis[7]+dis[8]是除了5的子树的其他点到4的距离和。这里不好求，如果我们求出所有点到4的距离s3那么就可以得到:s2=s3-s1$$

我们在向下转移到4节点时，s3的改变为4上面所有的点距离+1（边4-5）5包含自己的子孙节点个数-1。

那么就可以动态维护了。

#include<bits/stdc++.h>
#define LL long long
using namespace std;
const int mod=998244353;

vector<int> v[1000005];
LL siz[1000005];//子节点个数
LL sum[1000005];//子树到自己的距离和
LL d[1000005];//所有点到节点1的距离
LL s[1000005];//每个点到其他所有点的距离和
int n;
void dfs1(int u, int fa, int dis){//维护siz sum

    d[u]=dis;
    for(int i=0; i<v[u].size(); i++){
        int to=v[u][i];
        if(to!=fa){
            dfs1(to, u, dis+1);
            sum[u]+=(sum[to]+siz[to]);
            siz[u]+=siz[to];
        }
    }
    siz[u]+=1;
}

void dfs2(int u, int fa, LL s3){

    for(int i=0; i<v[u].size(); i++){
        int to=v[u][i];
        if(to!=fa){
            s[to]=s[u];
            LL s1=s3-sum[to]-siz[to];
            s1%=mod;
            LL s2=sum[to]+siz[to];
            s2%=mod;
            s[to]+=2*s1+n;
            s[to]-=2*s2;
            while(s[to]<0){
                s[to]=(s[to]+mod)%mod;
            }
            dfs2(to, u, s3+n-2*siz[to]);
        }
    }
}

int main()
{
    scanf("%d", &n);
    for(int i=1; i<n; i++){
        int a, b;
        scanf("%d%d", &a, &b);
        v[a].push_back(b);
        v[b].push_back(a);
    }
    dfs1(1, -1, 0);
    for(int i=2; i<=n; i++){//计算出1的值
        s[1]+=d[i]*d[i]%mod;
    }
    dfs2(1, -1, sum[1]);//动态转移
    LL ans=0;
    for(int i=1; i<=n; i++){
        ans=(ans+s[i])%mod;
    }
    cout<<ans<<endl;

    return 0;
}