牛客练习赛55 树（树形dp）

题目链接：点击这里

题目大意：
给一颗 $n$ 个节点的树，边权为 $1$ ，求 ${\sum }_{i=1}^n{\sum }_{j=1}^{n}di{s}^2(i,j)$ 对 $998244353$ 取模的结果

题目分析：
$$dis(i,j)=de{p}_i+de{p}_j-2de{p}_{lca(i,j)}$$
$$di{s}^2(i,j)=(de{p}_i+de{p}_j{)}^2-4de{p}_{lca(i,j)}(de{p}_i+de{p}_j)+4de{p}_{lca(i,j)}^2$$
$$=de{p}_i^2+de{p}_j^2+2de{p}_{i}de{p}_j-4de{p}_{lca(i,j)}(de{p}_i+de{p}_j)+4de{p}_{lca(i,j)}^2$$
可以发现前 $3$ 项是可以通过预处理各个节点深度后 $O(n)$ 求解，此时问题转化成了如何求后两项
我们可以转换思路，通过枚举 $lca(i,j)=x$ ，求出 $lca(i,j)=x$ 的对数和 $de{p}_i+de{p}_j$ 就可以求解出上式的答案了
我们定义 $si{z}_i$ 表示根为 $i$ 的子树的大小， $su{m}_i$ 表示根为 $i$ 的子树的深度之和
那么 $lca(i,j)=x$ 对数即为 $si{z}_x\ast si{z}_i-{\sum }_{y\in so{n}_x}si{z}_y\ast si{z}_y$ ，此式子可以理解为我们先假设 $x$ 与其子树的所有节点都相关联，再从中减去其所有子树中之和子树相关联的点的个数
这些点对的 $de{p}_i+de{p}_j$ 的和为 $2{\sum }_{y\in so{n}_x}su{m}_y\ast (si{z}_x-si{z}_y)+2de{p}_x\ast si{z}_x$
前半部分乘 $2$ 是因为 $(i,j),(j,i)$ 这样的点对是都要计数的，后半部分是点 $x$ 自身产生的贡献

具体细节见代码：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
#include<set>
#include<map>
#define ll long long
#define inf 0x3f3f3f3f
#define Inf 0x3f3f3f3f3f3f3f3f
#define int ll
using namespace std;
int read()
{
	int res = 0,flag = 1;
	char ch = getchar();
	while(ch<'0' || ch>'9')
	{
		if(ch == '-') flag = -1;
		ch = getchar();
	}
	while(ch>='0' && ch<='9')
	{
		res = (res<<3)+(res<<1)+(ch^48);//res*10+ch-'0';
		ch = getchar();
	}
	return res*flag;
}
const int maxn = 1e6+5;
const int maxm = 3e4+5;
const int mod = 998244353;
const double pi = acos(-1);
const double eps = 1e-8;
struct Edge{
	int nxt,to;
}edge[maxn<<1];
int n,cnt,ans,head[maxn],dep[maxn],siz[maxn],sum[maxn];
void addedge(int from,int to)
{
	edge[++cnt].nxt = head[from];
	edge[cnt].to = to;
	head[from] = cnt;
}
void dfs1(int u,int fa)
{
	siz[u] = 1;dep[u] = dep[fa]+1;sum[u] = dep[u];
	for(int i = head[u];i;i = edge[i].nxt)
	{
		int to = edge[i].to;
		if(to == fa) continue;
		dfs1(to,u);
		siz[u] += siz[to];
		sum[u] += sum[to];
	}
}
void dfs2(int u,int fa)
{
	int num = siz[u]*siz[u],res = 2*dep[u]*siz[u]%mod;
	for(int i = head[u];i;i = edge[i].nxt)
	{
		int to = edge[i].to;
		if(to == fa) continue;
		dfs2(to,u);
		num -= siz[to]*siz[to];
		res += 2*sum[to]*(siz[u]-siz[to])%mod;
	}
	ans = ((ans+4*num*dep[u]%mod*dep[u]%mod-4*dep[u]*res%mod)%mod+mod)%mod; 
}
signed main()
{
	n = read();
	for(int i = 1;i < n;i++)
	{
		int u = read(),v = read();
		addedge(u,v); addedge(v,u);
	} 
	dfs1(1,0);
	int tmp = 0;
	for(int i = 1;i <= n;i++) tmp = (tmp+dep[i])%mod,ans = (ans+dep[i]*dep[i]%mod*2*n%mod)%mod;
	ans = (ans+2*tmp*tmp%mod)%mod;
	dfs2(1,0);
	printf("%lld\n",ans);
	return 0;
}