Codeforces 981H：K Paths

传送门

考虑枚举一条路径 \(u,v\)，求出所有边经过它的答案

只需要求出 \(u\) 的子树内选出 \(k\) 个可以重复的点，使得它们到 \(u\) 的路径不相交

不难发现，就是从 \(u\) 的儿子的子树内各自选一个以及可以选多次 \(u\) 自己

设这个方案数为 \(f_u\)

再设 \(size_u\) 表示 \(u\) 的子树大小，\(son_u\) 表示 \(u\) 的儿子集合

考虑生成函数，设

\[A(x)=\prod_{v\in son_u}(1+size_vx)
\]

那么 \(f_u=\sum_{i=0}^{k}k^{\underline i}A_i\)

运用分治 \(FFT\) 预处理 \(f\) 可以做到 \(O(nlog^2n)\)

现在的问题是当 \(u,v\) 是祖先关系的时候，设 \(u\) 为 \(v\) 的祖先

那么 \(u\) 要计算以 \(v\) 为根的时候的子树答案

一个想法就是两遍 \(dfs\) 每次到一条边 \(u\rightarrow v\) 对于这个多项式乘上一个

\[\frac{1+(n-size_u)x}{1+size_vx}
\]

而考虑到一个点的儿子不同的 \(size\) 最多之后 \(\sqrt{n}\) 个，那么只要做 \(\sqrt{n}\) 次上面的操作，复杂度 \(O(n\sqrt{n})\)

总复杂度 \(O(nlog^2n+n\sqrt{n})\)

枚举路径可以通过子树和优化成 \(O(n)\)

# include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int maxn(4e5 + 5);
const int mod(998244353);

inline void Inc(int &x, int y) {
    x = x + y >= mod ? x + y - mod : x + y;
}

inline void Dec(int &x, int y) {
    x = x - y < 0 ? x - y + mod : x - y;
}

inline int Add(int x, int y) {
    return x + y >= mod ? x + y - mod : x + y;
}

inline int Sub(int x, int y) {
    return x - y < 0 ? x - y + mod : x - y;
}

inline int Pow(ll x, int y) {
    ll ret = 1;
    for (; y; y >>= 1, x = x * x % mod)
        if (y & 1) ret = ret * x % mod;
    return ret;
}

int w[2][maxn], r[maxn], l, deg;

inline void Init(int len) {
	int i, x, y;
	for (l = 0, deg = 1; deg < len; deg <<= 1) ++l;
	for (i = 0; i < deg; ++i) r[i] = (r[i >> 1] >> 1) | ((i & 1) << (l - 1));
	x = Pow(3, (mod - 1) / deg), y = Pow(x, mod - 2), w[0][0] = w[1][0] = 1;
	for (i = 1; i < deg; ++i) w[0][i] = (ll)w[0][i - 1] * x % mod, w[1][i] = (ll)w[1][i - 1] * y % mod;
}

inline void DFT(int *p, int opt) {
	int i, j, k, t, wn, x, y;
	for (i = 0; i < deg; ++i) if (r[i] < i) swap(p[r[i]], p[i]);
	for (i = 1; i < deg; i <<= 1)
		for (t = i << 1, j = 0; j < deg; j += t)
			for (k = 0; k < i; ++k) {
				wn = w[opt == -1][deg / t * k];
				x = p[j + k], y = (ll)p[j + k + i] * wn % mod;
				p[j + k] = Add(x, y), p[j + k + i] = Sub(x, y);
			}
	if (opt == -1) for (i = 0, wn = Pow(deg, mod - 2); i < deg; ++i) p[i] = (ll)p[i] * wn % mod;
}

int val[maxn], len, tmpa[20][maxn], tmpb[20][maxn], *a, *b;

void Solve(int l, int r, int *ret, int d) {
	if (l == r) {
		ret[0] = 1, ret[1] = val[l];
		return;
	}
	int mid = (l + r) >> 1, tmp = r - l + 2, i;
	Solve(l, mid, tmpa[d], d + 1), Solve(mid + 1, r, tmpb[d], d + 1);
	a = tmpa[d], b = tmpb[d], Init(tmp), DFT(a, 1), DFT(b, 1);
	for (i = 0; i < deg; ++i) a[i] = (ll)a[i] * b[i] % mod;
	DFT(a, -1);
	for (i = 0; i < tmp; ++i) ret[i] = a[i];
	for (i = 0; i < deg; ++i) a[i] = b[i] = 0;
}

int n, k, first[maxn], cnt, size[maxn], ans, tmp[maxn], dv[maxn];
int f[maxn], fac[maxn], rec[maxn], inv[maxn], g[maxn], h[maxn];

struct Edge {
	int to, next;
} edge[maxn];

inline void AddEdge(int u, int v) {
	edge[cnt] = (Edge){v, first[u]}, first[u] = cnt++;
	edge[cnt] = (Edge){u, first[v]}, first[v] = cnt++;
}

inline int Calc(int len) {
	int i, ret = 0, r = min(len, k);
	for (i = 0; i <= r; ++i) Inc(ret, (ll)tmp[i] * fac[k] % mod * inv[k - i] % mod);
	return ret;
}

inline void PolyDiv(int len, int v) {
	int i, inv, invv = Pow(v, mod - 2);
	for (i = 0; i <= len; ++i) dv[i] = tmp[i], tmp[i] = 0;
	for (i = len; i; --i)
		if (dv[i]) {
			tmp[i - 1] = inv = (ll)dv[i] * invv % mod;
			Dec(dv[i], (ll)inv * invv % mod), Dec(dv[i - 1], inv);
		}
}

inline void PolyMul(int len, int v) {
	int i;
	for (i = len; i; --i) Inc(tmp[i], (ll)tmp[i - 1] * v % mod);
}

inline int Cmp(int x, int y) {
	return size[x] < size[y];
}

void Dfs1(int u, int ff) {
	int e, v, i, j;
	size[u] = 1;
	for (e = first[u]; ~e; e = edge[e].next)
		if ((v = edge[e].to) ^ ff) Dfs1(v, u), Inc(g[u], g[v]), size[u] += size[v];
	len = 0;
	for (e = first[u]; ~e; e = edge[e].next)
		if ((v = edge[e].to) ^ ff) val[++len] = size[v];
	if (!len) f[u] = 1, Inc(g[u], 1);
	else {
		Solve(1, len, tmp, 0), f[u] = Calc(len), Inc(g[u], f[u]), len = 0;
		for (e = first[u]; ~e; e = edge[e].next)
			if ((v = edge[e].to) ^ ff) val[++len] = v;
		sort(val + 1, val + len + 1, Cmp);
		for (i = 0; i <= len; ++i) rec[i] = tmp[i];
		for (i = 1; i <= len; ++i) {
			if (size[val[i]] == size[val[i - 1]]) h[val[i]] = h[val[i - 1]];
			else {
				for (j = 0; j <= len; ++j) tmp[j] = rec[j];
				PolyDiv(len, size[val[i]]), PolyMul(len, n - size[u]);
				h[val[i]] = Calc(len);
			}
		}
	}
}

void Dfs2(int u, int ff) {
	int e, v;
	for (e = first[u]; ~e; e = edge[e].next)
		if ((v = edge[e].to) ^ ff) Inc(ans, (ll)Sub(h[v], f[u]) * g[v] % mod), Dfs2(v, u);
}

int main() {
	memset(first, -1, sizeof(first));
	int i, u, v, sum = 0, mx;
	scanf("%d%d", &n, &k);
	if (k == 1) return printf("%lld\n", ((ll)n * (n - 1) >> 1) % mod), 0;
	mx = max(n, k);
	for (i = 1; i < n; ++i) scanf("%d%d", &u, &v), AddEdge(u, v);
	inv[0] = inv[1] = 1;
	for (i = 2; i <= mx; ++i) inv[i] = (ll)(mod - mod / i) * inv[mod % i] % mod;
	for (i = 2; i <= mx; ++i) inv[i] = (ll)inv[i - 1] * inv[i] % mod;
	for (i = fac[0] = 1; i <= mx; ++i) fac[i] = (ll)fac[i - 1] * i % mod;
	Dfs1(1, 0), Dfs2(1, 0);
	for (i = 1; i <= n; ++i) Inc(sum, f[i]);
	sum = (ll)sum * sum % mod;
	for (i = 1; i <= n; ++i) Dec(sum, (ll)f[i] * f[i] % mod);
	sum = (ll)sum * inv[2] % mod;
	Inc(ans, sum), printf("%d\n", ans);
	return 0;
}