题解 CF1770E【Koxia and Tree】

先对不好处理的部分做一个转化：任取两个蝴蝶并计算它们之间的距离，等价于所有蝴蝶两两距离的平均值。于是只需要求所有蝴蝶两两距离的和。

首先令每条边都按输入顺序定向，DFS 预处理出每条边两侧蝴蝶的数量，容易计算出初始局面下的答案。

然后按照输入顺序枚举每条边，也就是第二步的按顺序操作。我们维护 $dp$ 数组表示此时每个点有蝴蝶的概率，然后分别考虑 $u_i\to v_i$ 和 $v_i\to u_i$ 两种可能的蝴蝶移动，通过数学推导可以计算出此时每条边对答案的贡献以及新的概率。

式子不难推，但是比较长不方便列出，详见代码。

// Problem: E. Koxia and Tree
// Contest: Codeforces - Good Bye 2022: 2023 is NEAR
// URL: https://codeforces.com/contest/1770/problem/E
// Memory Limit: 256 MB
// Time Limit: 3000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

//By: OIer rui_er
#include <bits/stdc++.h>
#define rep(x,y,z) for(ll x=(y);x<=(z);x++)
#define per(x,y,z) for(ll x=(y);x>=(z);x--)
#define debug(format...) fprintf(stderr, format)
#define fileIO(s) do{freopen(s".in","r",stdin);freopen(s".out","w",stdout);}while(false)
using namespace std;
typedef long long ll;
const ll N = 3e5+5, mod = 998244353, inv2 = 499122177;

ll n, k, a[N], U[N], V[N], vis[N], dp[N], cntV[N], cntE[N];
vector<tuple<ll, ll>> e[N];
template<typename T> void chkmin(T& x, T y) {if(x > y) x = y;}
template<typename T> void chkmax(T& x, T y) {if(x < y) x = y;}
ll qpow(ll x, ll y) {
	ll ans = 1;
	for(; y; y >>= 1, x = x * x % mod) if(y & 1) ans = ans * x % mod;
	return ans;
}
ll inv(ll x) {return qpow(x, mod-2);}
void dfs(ll u, ll f) {
	cntV[u] = vis[u];
	for(auto i : e[u]) {
		ll v = get<0>(i), id = get<1>(i);
		if(v == f) continue;
		dfs(v, u);
		cntV[u] += cntV[v];
		cntE[id] = v == U[id] ? cntV[v] : k - cntV[v];
	}
}

int main() {
	scanf("%lld%lld", &n, &k);
	rep(i, 1, k) scanf("%lld", &a[i]);
	rep(i, 1, n-1) {
		scanf("%lld%lld", &U[i], &V[i]);
		e[U[i]].emplace_back(V[i], i);
		e[V[i]].emplace_back(U[i], i);
	}
	rep(i, 1, k) vis[a[i]] = 1;
	dfs(1, 0);
	ll ans = 0;
	rep(i, 1, n-1) ans = (ans + cntE[i] * (k - cntE[i]) % mod) % mod;
	rep(i, 1, n) dp[i] = vis[i];
	rep(i, 1, n-1) {
		ans = (ans + dp[U[i]] * (1 + mod - dp[V[i]]) % mod * (((cntE[i] - 1) - (k - cntE[i]) + 3 * mod) % mod) % mod * inv2 % mod) % mod;
		ans = (ans + dp[V[i]] * (1 + mod - dp[U[i]]) % mod * (((k - cntE[i] - 1) - (cntE[i]) + 3 * mod) % mod) % mod * inv2 % mod) % mod;
		ll x = dp[U[i]], y = dp[V[i]];
		dp[U[i]] = (x * inv2 % mod + x * y % mod * inv2 % mod + y * (1 + mod - x) % mod * inv2 % mod) % mod;
		dp[V[i]] = (y * inv2 % mod + x * y % mod * inv2 % mod + x * (1 + mod - y) % mod * inv2 % mod) % mod;
	}
	ans = ans * inv(k*(k-1)/2%mod) % mod;
	printf("%lld\n", ans);
	return 0;
}