P10842 【MX-J2-T3】Piggy and Trees 题解

P10842 【MX-J2-T3】Piggy and Trees 题解

solution

首先我们发现 \(f(u,v,i)\) 函数很难讨论，把它转换成 \(i\) 到路径 \(u \to v\) 的距离。这样在我们以 \(i\) 为根时就有 \(f(u,v,i)=dis(i,lca(u,v))\)。到了这里可以想到换根。我们定换根时初始的根为 \(root\)，统计当前点的子树中有多少个 \(lca(u,v)\) 以及对应的贡献 \(dis(i,lca(u,v))\) 的总和，我们设 \(x\) 的子树中 \(lca(u,v)\) 的个数为 \(d_x\)，子树中的贡献和为 \(s_x\)。这样对于每一个 \(s_x\) 我们都从 \(x\) 的子节点转移来，再因为转移时子节点的子树的根节点往上移了一层（也就是从 \(x\) 的子节点转移到了 \(x\)），所以每一个 \(lca(u,v)\) 的贡献都加上了 \(1\)，所以加上子树中 \(lca(u,v)\) 的个数。这样就能够算出换根的初始值，也就是 \(s_{root}\)。

接下来考虑如何转移。发现从 \(fa_x\) 转移至 \(x\) 会将 \(x\) 的子树内的贡献都减一，将 \(x\) 子树外的贡献都加一，于是我们直接转移。但是我们发现会算重，因为有 \(f(a,b,root)\) （其中 \(a\) 为 \(x\) 的子树内节点，\(b\) 为 \(x\) 的子树外节点）没有被减掉，这些部分我们单独计算即可。

code

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N = 2e5 + 10, mod = 1e9 + 7;
int n, s[N], d[N], r[N], siz[N], f[N], ans = 0;
vector<int> g[N];
void dfs(int u, int fa) {
	siz[u] = 1;
	int res = 0, sum = 0;
	for(auto v : g[u]) {
		if(v != fa) {
			dfs(v, u);
			s[u] += s[v] + d[v];
			d[u] += d[v];
			siz[u] += siz[v];
			res += siz[v] * sum;
			sum += siz[v];
		}
	}
	d[u] += res + siz[u] - 1;
	r[u] = (n - siz[u]) * siz[u];
}
void dp(int u, int fa) {
	f[u] = f[fa] - d[u] - r[u] + d[1] - d[u];
	if(!fa) f[u] = s[u];
	ans += f[u]; ans %= mod;
	for(auto v : g[u]) {
		if(v != fa) dp(v, u);
	}
}
signed main() {
	cin >> n;
	for(int i = 1, u, v; i < n; i++) {
		cin >> u >> v;
		g[u].push_back(v); g[v].push_back(u);
	}
	dfs(1, 0);
	dp(1, 0);
	cout << ans;
	return 0;
}