Cattle-off practice match 55 E- tree tree DP

The meaning of problems

Do you have a size of \ (n \) tree, from the point \ (1 \) to \ (n \) label.
Set \ (dis⁡ (x, y) \) represents \ (X \) to \ (Y \) distance.
Seeking \ (\ sum_ {i = 1 } ^ {n} \ sum_ {j = 1} ^ {n} dis ^ 2 (i, j) \) to (998244353 \) \ result of modulo.

analysis

Set \ (d_x \) is the point \ (x \) depth.

\(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 vigorously tree to get away.

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;
}