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