「SOL」数树 (LOJ/WC2019)

WC 果然还是 WC

---

# 题面

有一张 $n$ 个点的图，图上有红蓝两种边（可能重叠），且两种边各自形成一个 $n$ 个点的树。

用 $m$ 种颜色给图上的所有点染色。若 $u,v$ 两点在红边上的路径和在蓝边上的路径完全重合，则 $u,v$ 必须染相同的颜色。

求解下列三种问题：

1. 给定所有红边和蓝边，求染色方案数；
2. 给定所有蓝边，对红边的所有可能情况，求染色方案数之和；
3. 不给定边，对蓝边和红边的所有可能情况，求染色方案数之和。

数据规模：$n\le10^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$，则产生 $\times a_i^{d_i}$ 的贡献。于是图 $G$ 对应的生成树 $T$ 的方案数为：

$$ \begin{aligned} &\frac{(m-2)!}{\prod(d_i-1)!}\prod a_i^{d_i}\\ =&\prod a_i\times\frac{(m-2)!}{\prod(d_i-1)!}\prod a_i^{d_i-1} \end{aligned} $$

对于所有 $\{d_i\}$ 的情况求和，注意到 $\sum (d_i-1)=m-2$，于是所有生成树的答案为：

$$ \begin{aligned} &(m-2)!\prod a_i\times\left(\sum_{\sum(d_i-1)=m-2}\frac{a_i^{d_i-1}}{(d_i-1)!}\right)\\ =&(m-2)!\prod a_i\times [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_{j\ge i}\binom{j}{i}g_j\Rightarrow g_i=\sum_{j\ge i}\binom{j}{i}(-1)^{j-i}f_j$$

这样直接做是 $\mathcal{O}(n^2)$ 的，需要再推式子。答案的式子为：

$$\begin{aligned} &\sum_{i=0}^{n-1}g_im^{n-i}\\ =&\sum_{i=0}^{n-1}\sum_{j\ge i}\binom{j}{i}(-1)^{j-i}f_jm^{n-i}\\ =&\sum_{j=0}^{n-1}f_jm^n\sum_{i=0}^{j}\binom{j}{i}(-1)^{j-i}m^{-i}\\ =&m^n\sum_{j=0}^{n-1}f_j(m^{-1}-1)^{j} \end{aligned}$$

这个式子就非常有用了——每钦定一条边是双色边，就会产生 $\times(m^{-1}-1)$ 的贡献。于是就不需要进行树形背包，省去一维，复杂度变为 $\mathcal{O}(n)$。

- 问题 3

仿照问题 2 的思路，不过现在没有蓝色边的限制，我们可以直接决策连通块的大小。

可以设计出 DP：将 $1\sim i$ 划分成若干连通块，对答案的贡献之和。转移考虑枚举 $\mathbf{1}$ 所在的连通块大小，注意现在钦定一条边是双色边会影响两棵树，相应式子需要平方：

$$\begin{aligned} f_i& =\sum_{j=1}^{i}\binom{i-1}{j-1}f_{i-j}(n^{-2}(m^{-1}-1))^{j-1}\\ \Rightarrow \frac{f_i}{(i-1)!}& =\sum_{j=1}^i\frac{f_{i-j}}{(i-j)!}\frac{(n^{-2}(m^{-1}-1))^{j-1}}{(j-1)!} \end{aligned}$$

这是一个分治FFT的式子（CDQ分治），可以做到 $\mathcal{O}(n\log^2n)$，但是这还不够。

考虑 $f_i$ 的生成函数 $F(x)$。观察等式左侧 $\frac{f_i}{(i-1)!}=\frac{f_i}{i!}\times i$，对应 $F’(x)$。

记

$$G(x)=\sum_{i\ge0}\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 G\Rightarrow 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\n", 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\n", 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\n", 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\n", 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;
}

---

THE END

Thanks for reading!

你就是泪浸白了初雪
才会如离人来去飘洒摇曳
你就是你染红了岁月
改变我黑白而无言的世界

——《寄明月(Cover)》By 祖娅纳惜

> Link 寄明月 - Bilibili