WC 果然还是 WC
# 题面
有一张 (n) 个点的图，图上有红蓝两种边（可能重叠），且两种边各自形成一个 (n) 个点的树。
用 (m) 种颜色给图上的所有点染色。若 (u,v) 两点在红边上的路径和在蓝边上的路径完全重合，则 (u,v) 必须染相同的颜色。
求解下列三种问题：
- 给定所有红边和蓝边，求染色方案数；
- 给定所有蓝边，对红边的所有可能情况，求染色方案数之和；
- 不给定边，对蓝边和红边的所有可能情况，求染色方案数之和。
数据规模：(nle10^5)。
# 解析
(u,v) 必须染相同颜色的条件其实就是「只保留红蓝重叠的双色边，(u,v) 在相同连通块中」。于是答案就是 (m) 的“连通块个数”次方。
- 问题 1
直接计算上述连通块数量。
- 问题 2
对连通块状态进行 DP，只需要决策哪些边是双色边，进而判断出连通块数量。
首先根据 prufer 序列，有结论：
结论
一个 $n$ 个点的有标号图 $G$，已有的边形成了 $m$ 棵树的森林，每棵树的大小依次为 $a_1,a_2,dots a_m$，随意加边形成一棵树的方案数为
$$ n^{m-2}prod_{i=1}^m a_i $$证明如下：将已有的连通块缩点得到图 $G_0$（第 $i$ 个连通块对应点 $i$）。若 $G_0$ 的生成树 $T_0$ 满足点 $i$ 的度数为 $d_i$，根据 prufer 序列的性质，$G_0$ 的 prufer 序列中，$i$ 会出现 $d_i-1$ 次。用可重排的方法分配位置，满足条件的生成树 $T_0$ 的方案数为：
$$ frac{(m-2)!}{prod (d_i-1)!} $$还原到图 $G$，$T_0$ 中每条边的端点都可以在它对应的连通块内任取，第 $i$ 个连通块度数为 $d_i$，则产生 $ imes a_i^{d_i}$ 的贡献。于是图 $G$ 对应的生成树 $T$ 的方案数为：
$$ egin{aligned} &frac{(m-2)!}{prod(d_i-1)!}prod a_i^{d_i}\ =&prod a_i imesfrac{(m-2)!}{prod(d_i-1)!}prod a_i^{d_i-1} end{aligned} $$对于所有 ${d_i}$ 的情况求和，注意到 $sum (d_i-1)=m-2$，于是所有生成树的答案为：
$$ egin{aligned} &(m-2)!prod a_i imesleft(sum_{sum(d_i-1)=m-2}frac{a_i^{d_i-1}}{(d_i-1)!} ight)\ =&(m-2)!prod a_i imes [x^{m-2}]exp(Sigma a_ix) end{aligned} $$显然 $sum a_i=n$，于是上式即为：
$$ (m-2)![x^{m-2}]exp(nx)prod a_i=n^{m-2}prod a_i $$上述结论的 (prod a_i) 可以处理成“每个连通块选取一个点”，而 (n^{m-2}) 可以处理成“初始值为 (n^{n-2})，每连接一条双色边，就少一个连通块，答案乘上 (n^{-1})”。
于是这样可以得到一个树形背包+二项式反演的思路，(f(u,i,0/1)) 表示钦定 (u) 子树内有 (i) 条双色边，且 (u) 所在的连通块 有/没有 选点的生成树个数；转移就讨论 ((u,v)) 是否是红蓝重叠的边，如果是，则连通块个数减少 (1)，方案数乘上 (n^{-1})。
记 (f_i) 表示「钦定 (i) 条边是双色边的生成树个数」。要得到「恰好 (i) 条边」，是一个经典的二项式反演，记 (g_i) 表示「恰有 (i) 条双色边的生成树个数」：
[f_i=sum_{jge i}inom{j}{i}g_jRightarrow g_i=sum_{jge i}inom{j}{i}(-1)^{j-i}f_j
]
这样直接做是 (mathcal{O}(n^2)) 的，需要再推式子。答案的式子为：
[egin{aligned}
&sum_{i=0}^{n-1}g_im^{n-i}\
=&sum_{i=0}^{n-1}sum_{jge i}inom{j}{i}(-1)^{j-i}f_jm^{n-i}\
=&sum_{j=0}^{n-1}f_jm^nsum_{i=0}^{j}inom{j}{i}(-1)^{j-i}m^{-i}\
=&m^nsum_{j=0}^{n-1}f_j(m^{-1}-1)^{j}
end{aligned}
]
这个式子就非常有用了——每钦定一条边是双色边，就会产生 ( imes(m^{-1}-1)) 的贡献。于是就不需要进行树形背包，省去一维，复杂度变为 (mathcal{O}(n))。
- 问题 3
仿照问题 2 的思路，不过现在没有蓝色边的限制，我们可以直接决策连通块的大小。
可以设计出 DP：将 (1sim i) 划分成若干连通块，对答案的贡献之和。转移考虑枚举 (mathbf{1}) 所在的连通块大小，注意现在钦定一条边是双色边会影响两棵树，相应式子需要平方：
[egin{aligned}
f_i& =sum_{j=1}^{i}inom{i-1}{j-1}f_{i-j}(n^{-2}(m^{-1}-1))^{j-1}\
Rightarrow frac{f_i}{(i-1)!}& =sum_{j=1}^ifrac{f_{i-j}}{(i-j)!}frac{(n^{-2}(m^{-1}-1))^{j-1}}{(j-1)!}
end{aligned}
]
这是一个分治FFT的式子（CDQ分治），可以做到 (mathcal{O}(nlog^2n))，但是这还不够。
考虑 (f_i) 的生成函数 (F(x))。观察等式左侧 (frac{f_i}{(i-1)!}=frac{f_i}{i!} imes i)，对应 (F'(x))。
记
[G(x)=sum_{ige0}frac{(n^{-2}(m^{-1}-1))^{i}}{i!}x^i
]
则有 (F'=FG)，即 (frac{F'}{F}=G)。而 (frac{F'}{F}) 是 (ln(F)) 对 (mathbf{x}) 求导的结果，于是两边对 (mathbf{x}) 积分：
[ln(F)=int GRightarrow F=e^{int G}
]
可以直接多项式 Exp 计算。
# 源代码
三合一的代码太长了……
点击展开/折叠代码
/*Lucky_Glass*/
#include <set>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
inline int rin(int &r) {
int b = 1, c = getchar(); r = 0;
while ( c < '0' || '9' < c ) b = c == '-' ? -1 : b, c = getchar();
while ( '0' <= c && c <= '9' ) r = (r * 10) + (c ^ '0'), c = getchar();
return r *= b;
}
const int N = 1e5 + 10, MOD = 998244353, L = 1 << 18;
#define con(typ) const typ &
inline int add(int a, con(int) b) {return (a += b) >= MOD ? a - MOD : a;}
inline int sub(int a, con(int) b) {return (a -= b) < 0 ? a + MOD : a;}
inline int mul(con(int) a, con(int) b) {return int(1ll * a * b % MOD);}
inline int iPow(int a, int b) {
int r = 1;
while ( b ) {
if ( b & 1 ) r = mul(r, a);
a = mul(a, a), b >>= 1;
}
return r;
}
int n, m, datatyp;
int fac[L + 2], ifac[L + 2];
inline int binom(con(int) u, con(int) v) {
return u >= v ? mul(fac[u], mul(ifac[v], ifac[u - v])) : 0;
}
namespace SUBTASK0 {
typedef std::pair<int, int> pii;
std::set<pii> treb;
void solve() {
for (int i = 1, u, v; i < n; ++i) {
rin(u), rin(v);
if ( u > v ) std::swap(u, v);
treb.insert(std::make_pair(u, v));
}
int cnt = n;
for (int i = 1, u, v; i < n; ++i) {
rin(u), rin(v);
if ( u > v ) std::swap(u, v);
cnt -= treb.count(std::make_pair(u, v));
}
printf("%d
", iPow(m, cnt));
}
}
namespace SUBTASK1 {
struct Graph {
int head[N], nxt[N << 1], to[N << 1], ncnt;
void addEdge(con(int) u, con(int) v) {
int p = ++ncnt, q = ++ncnt;
to[p] = v, nxt[p] = head[u], head[u] = p;
to[q] = u, nxt[q] = head[v], head[v] = q;
}
inline int operator [] (con(int) u) const {return head[u];}
} gr;
int invn, cony;
int f[N][2];
void dfs(con(int) u, con(int) fa) {
f[u][0] = f[u][1] = 1;
for (int it = gr[u]; it; it = gr.nxt[it]) {
int v = gr.to[it];
if ( v == fa ) continue;
dfs(v, u);
int fu0 = f[u][0], fu1 = f[u][1];
f[u][0] = add(
mul(mul(fu0, f[v][0]), cony), // link (u, v)
mul(fu0, f[v][1])
);
f[u][1] = add(mul(fu1, f[v][1]), add(
mul(mul(fu0, f[v][1]), cony),
mul(mul(fu1, f[v][0]), cony)
));
}
}
void solve() {
for (int i = 1, u, v; i < n; ++i) gr.addEdge(rin(u), rin(v));
invn = iPow(n, MOD - 2), cony = sub(iPow(m, MOD - 2), 1);
cony = mul(cony, invn);
dfs(1, 0);
printf("%d
", mul(mul(iPow(n, n - 2), iPow(m, n)), f[1][1]));
}
}
namespace SUBTASK2 {
namespace BASICPOLY {
int rev[L], elg2[L + 2], powg[L + 2], invi[L + 2];
void init() {
elg2[1] = 0, powg[0] = 1, powg[1] = iPow(3, (MOD - 1) >> 18), invi[1] = 1;
for (int i = 2; i <= L; ++i) {
elg2[i] = elg2[(i + 1) >> 1] + 1;
powg[i] = mul(powg[i - 1], powg[1]);
invi[i] = mul(ifac[i], fac[i - 1]);
}
}
void ntt(int *arr, con(int) len, con(int) typ) {
for (int i = 1; i < len; ++i) {
rev[i] = rev[i >> 1] >> 1 | ((i & 1) ? (len >> 1) : 0);
if ( rev[i] < i ) std::swap(arr[i], arr[rev[i]]);
}
for (int i = 1, ii = 2; i < len; i <<= 1, ii <<= 1) {
int s = L >> elg2[ii];
for (int j = 0; j < len; j += ii) {
int *a = arr + j, *b = a + i, *p = powg, q = *b;
for (int k = 0; k < i; ++k, ++a, q = mul(*(++b), *(p += s)))
*b = sub(*a, q), *a = add(*a, q);
}
}
if ( typ == -1 ) {
std::reverse(arr + 1, arr + len);
for (int i = 0, ivn = MOD - ((MOD - 1) >> elg2[len]); i < len; ++i)
arr[i] = mul(arr[i], ivn);
}
}
int arr1[L], arr2[L], arr3[L], arr4[L];
void polyInv(int *a, int *b, con(int) len) {
if ( len == 1 ) {b[0] = iPow(a[0], MOD - 2); return;}
polyInv(a, b, (len + 1) >> 1);
int llen = 1 << elg2[len << 1];
for (int i = 0; i < len; ++i) arr1[i] = a[i];
for (int i = len; i < llen; ++i) arr1[i] = 0;
for (int i = 0, ii = (len + 1) >> 1; i < ii; ++i) arr2[i] = b[i];
for (int i = (len + 1) >> 1; i < llen; ++i) arr2[i] = 0;
ntt(arr1, llen, 1), ntt(arr2, llen, 1);
for (int i = 0; i < llen; ++i)
arr1[i] = mul(arr2[i], sub(2, mul(arr1[i], arr2[i])));
ntt(arr1, llen, -1);
for (int i = 0; i < len; ++i) b[i] = arr1[i];
}
void polyDer(int *a, int *b, con(int) len) {
for (int i = 0; i < len - 1; ++i)
b[i] = mul(a[i + 1], i + 1);
b[len - 1] = 0;
}
void polyInt(int *a, int *b, con(int) len) {
for (int i = len - 1; i; --i)
b[i] = mul(a[i - 1], invi[i]);
b[0] = 0;
}
void polyLn(int *a, int *b, con(int) len) {
int llen = 1 << elg2[len << 1];
polyInv(a, arr3, len);
for (int i = len; i < llen; ++i) arr3[i] = 0;
polyDer(a, arr1, len);
for (int i = len; i < llen; ++i) arr1[i] = 0;
ntt(arr3, llen, 1), ntt(arr1, llen, 1);
for (int i = 0; i < llen; ++i) arr3[i] = mul(arr3[i], arr1[i]);
ntt(arr3, llen, -1);
polyInt(arr3, arr3, len);
for (int i = 0; i < len; ++i) b[i] = arr3[i];
}
void polyExp(int *a, int *b, con(int) len) {
if ( len == 1 ) {b[0] = 1; return;}
polyExp(a, b, (len + 1) >> 1);
int llen = 1 << elg2[len << 1];
for (int i = (len + 1) >> 1; i < len; ++i) b[i] = 0;
polyLn(b, arr4, len);
for (int i = 0; i < len; ++i) arr4[i] = sub(a[i], arr4[i]);
arr4[0] = add(arr4[0], 1);
for (int i = len; i < llen; ++i) arr4[i] = 0;
for (int i = 0; i < len; ++i) arr1[i] = b[i];
for (int i = len; i < llen; ++i) arr1[i] = 0;
ntt(arr1, llen, 1), ntt(arr4, llen, 1);
for (int i = 0; i < llen; ++i) arr1[i] = mul(arr1[i], arr4[i]);
ntt(arr1, llen, -1);
for (int i = 0; i < len; ++i) b[i] = arr1[i];
}
}
int arr[N], arr2[N], res[N];
void db() {
using namespace BASICPOLY;
int a[20] = {0, 1}, b[20] = {};
polyExp(a, b, 4);
for (int i = 0; i < 16; ++i) printf("%d
", mul(b[i], fac[i]));
exit(0);
}
void solve() {
BASICPOLY::init();
int per = iPow(n, MOD - 2);
per = mul(mul(per, per), sub(iPow(m, MOD - 2), 1));
for (int i = 1, tmp = 1; i <= n; ++i, tmp = mul(tmp, per))
arr[i - 1] = mul(iPow(i, i), mul(tmp, ifac[i - 1]));
BASICPOLY::polyInt(arr, arr, n + 1);
BASICPOLY::polyExp(arr, res, n + 1);
printf("%d
", mul(
mul(res[n], fac[n]),
mul(iPow(m, n), iPow(n, 2 * n - 4))
));
}
}
void init() {
fac[0] = 1;
for (int i = 1; i <= L; ++i) fac[i] = mul(fac[i - 1], i);
ifac[L] = iPow(fac[L], MOD - 2);
for (int i = L - 1; ~i; --i) ifac[i] = mul(ifac[i + 1], i + 1);
}
int main() {
init();
rin(n), rin(m), rin(datatyp);
if ( datatyp == 0 ) SUBTASK0::solve();
else if ( datatyp == 1 ) SUBTASK1::solve();
else SUBTASK2::solve();
return 0;
}