题意
你有一颗大小为\(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;
}