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;

}

巴特西

Codeforces 981H：K Paths

最新文章

热门文章