牛客练习赛55 E-树 树形DP

题意

你有一颗大小为\(n\)的树，点从\(1\)到\(n\)标号。

设\(dis⁡(x,y)\)表示\(x\)到\(y\)的距离。

求\(\sum_{i=1}^{n}\sum_{j=1}^{n}dis^2(i,j)\)对\(998244353\)取模的结果。

分析

设\(d_x\)为点\(x\)的深度。

\(dis(x,y)=d_x+d_y-2*lca(x,y)\)

\(dis^2(x,y)=d_x^2+d_y^2+2*d_x*d_y-4*d_{lca(x,y)}*(d_x+d_y)+4*d_{lca(x,y)}^2\)

大力树dp就完事了。

Code

#include<bits/stdc++.h>
#define fi first
#define se second
#define lson l,mid,p<<1
#define rson mid+1,r,p<<1|1
#define pb push_back
#define ll long long
using namespace std;
const int inf=1e9;
const int mod=998244353;
const int maxn=1e6+10;
int n;
ll sz[maxn],d[maxn],dep[maxn];
vector<int>g[maxn];
ll ans;
void dfs(int u,int fa){
    sz[u]=1;d[u]=d[fa]+1;
    dep[u]=d[u];
    for(int x:g[u]){
        if(x==fa) continue;
        dfs(x,u);
        sz[u]=(sz[u]+sz[x])%mod;
        dep[u]=(dep[u]+dep[x])%mod;
    }
    for(int x:g[u]){
        if(x==fa) continue;
        ans=(ans+sz[x]*(sz[u]-sz[x])%mod*4%mod*d[u]%mod*d[u]%mod)%mod;
        ans=(ans+4*d[u]%mod*d[u]%mod*sz[x]%mod)%mod;
        ans=(ans-dep[x]*(sz[u]-sz[x])%mod*8%mod*d[u]%mod)%mod;
        ans=(ans-d[u]*sz[x]%mod*8%mod*d[u]%mod)%mod;
    }
}
int main(){
    //ios::sync_with_stdio(false);
    //freopen("in","r",stdin);
    scanf("%d",&n);
    for(int i=1,x,y;i<n;i++){
        scanf("%d%d",&x,&y);
        g[x].pb(y);g[y].pb(x);
    }
    dfs(1,0);
    for(int i=1;i<=n;i++){
        ans=(ans+d[i]*d[i]%mod*2%mod*(n-1)%mod)%mod;
        ans=(ans+d[i]*(dep[1]-d[i])%mod*2%mod)%mod;
    }
    ans=(ans+mod)%mod;
    printf("%lld\n",ans);
    return 0;
}