P8202

· · 题解

G. 染色

题意：给定一棵 n 个节点的树，要求用 m 种颜色将其染色，要求相邻两个节点的颜色不同，且任何一种颜色的使用次数都不超过 \lfloor\frac{n}{3}\rfloor + 1。

**关键词**：容斥，DP。 **参考难度**：蓝/紫。 **解析**：本题出题灵感来源于 「CF1487G String Counting」和「LuoguP5664 [CSP-S2019] Emiya 家今天的饭」。 首先注意到 $\lfloor\frac{n}{3}\rfloor + 1$ 的限制，可以发现即使我们不考虑这个限制，也至多会有两种颜色超出（因为一旦有三种颜色超出则总数大于 $n$）。于是不难设计容斥：答案 = 不考虑限制的方案数 - 至少有一种颜色超出限制的方案数 + 两种颜色超出限制的方案数。 我们直接考虑计算两种颜色超出限制的方案数：不难设计 dp：设 $f_{u, p,q,i, j, 0/1/2}$ 表示 $p$ 和 $q$ 是我们钦定的两种**可能**超限制的颜色，以 $u$ 为根的子树里有 $i$ 个 $p$ 颜色和 $j$ 个 $q$ 颜色，且 $u$ 是颜色 $p$ /颜色 $q$ /其他颜色的方案数，转移暂时不表。 进一步考虑发现颜色之间没有本质区别也就是对于不同的 $p, q$，算出 $f_{u, p, q, i, j, 0/1/ 2}$ 的结果都是一样的，于是我们只需要计算一次 $f$ 数组，然后乘上 $C_m^2$ 即可。于是状态被优化为：设 $f_{u, i, j, 0/1/2}$ 表示以 $u$ 为根的子树里有 $i$ 个颜色 1 和 $j$ 个颜色 2，且 $u$ 是颜色 1/颜色 2/其他颜色的方案数。 方便起见，我们设已经转移了 $u$ 的前 $p - 1$ 个孩子，目前考虑合并第 $p$ 个孩子 $v$ 的信息。已合并的孩子的信息存储在 $h_{p, i, j, 0/1/2}$ 数组中。 $$h_{p, i, j, 0} = \sum_{x = 0}^i\sum_{y = 0}^j(f_{v,x,y,0} + f_{v,x,y,1}) \times h_{p,i - x,j - y,0}$$ $h_{p, i, j, 1}$ 的转移与之完全对称。 考虑 $h_{p, i, j, 2}$：此时 $u$ 是一个非颜色 1 或颜色 2 的其他颜色。$v$ 是颜色 1/2 的转移是简单的，考虑 $v$ 也是一个其他颜色时，$v$ 的颜色不能和 $u$ 的相同。因此 $v$ 对 $h_p$ 的贡献不是 $f_{v, ?,?,2}$ 全部，而是除去了 $v$ 的颜色与 $u$ 相同的那部分。注意到 $f_{v}$ 共表示了 $v$ 取 $(m - 2)$ 种颜色的方案数，而颜色之间没有本质区别，因此每种颜色的方案数都是 $\frac {f_{v, ?, ?, 2}}{m - 2}$。去掉 $u$ 的颜色，共有 $m - 3$ 种颜色产生了贡献，因此总的贡献是 $\frac{(m - 3)f_{v, ?, ?, 2}}{m - 2}$。 于是转移如下： $h_{p, i, j, 2} = \sum_{x= 0}^i\sum_{y = 0}^j(f_{v,x,y,0} + f_{v,x,y,1} + \frac{(m - 3)f_{v,x,y,2}}{m - 2}) \times h_{p - 1, i - x, j - y, 2}

因为是模意义下的计数，所以除以 m - 2 要改为乘 m - 2 的逆。但是如果将 f_{u, ?, ?, 2} 的方案数定义为某一个不会超出限制的“颜色 3”的方案数，则可以规避掉乘逆元的过程。

上述转移的边界条件是 h_{p, 1, 0, 0} = h_{p, 0, 1, 1} = 1，h_{p,0,0,2} = m - 2，其余均为 0。若想要规避求逆元，则令 h_{p, 0, 0, 2} = 1 即可，最后计数时要把 (m - 2) 再乘上。

不难发现，\sum_{i = lim + 1}^n\sum_{j = 1}^nf_{1, i, j, 0/1/2} 即为至少有一个超出限制的方案数，其中 lim = \lfloor\frac{n}{3}\rfloor + 1。

将 i 改为从 1 至 n 枚举，就是不考虑限制的方案数。

当然，不考虑限制的方案数也可以这样推导：

设 g_{u, i} 表示以 u 为根的子树，u 的颜色是 i 的方案数；G_u 表示以 u 为根的子树的总方案数，则：

G_{u} = \sum\limits_{col = 1}^m\prod\limits_{v \in son(u)} \sum\limits_{1 \leq j \leq m}^{j \neq col}g_{v,j} = \sum\limits_{col = 1}^m \prod\limits_{v \in son(u)} \frac{m - 1}{m}G_v = m\prod\limits_{v \in son(u)} \frac{m - 1}{m}G_v

这里同样用到了颜色没有本质不同。

考虑 f 有 O(n^3) 个状态，转移是 O(n^2) 的，所以看起来的时间复杂度是 O(n^5)。但是可以使用如下的技巧将其实际复杂度优化至 O(n^4)：

void dfs(const int u, const int f) {
  sz[u] = 1;
  for (auto v : e[u]) if (v != f) {
    dfs(v);
    for (int i = 1; i <= sz[u]; ++i) {
      for (int x = 1; x <= sz[v]; ++x) {
        for (int j = 1; j <= sz[u]; ++j) {
          for (int y = 1; y <= sz[v]; ++y) {
            // do sth O(1).
          }
        }
      }
    }
    sz[u] += sz[v];
  }
}

这是因为枚举 i 和 x 的过程可以看作枚举 u 前面的子树和 v 这个子树里面的点组成的点对，显然整棵树的点对只有 O(n^2) 个，每个点对只会在 LCA 处被枚举。所以遍历树上 n 个点并枚举子树间的点对的总复杂度是 O(n^2)，而里面两层循环也是 O(n^2)，故总复杂度是 O(n^4)。

如果因为常数极小，O(n^5) 的算法有可能卡过本题。在测试中，五次方算法的最坏用时在 800ms 左右，因此本题的时间限制开到了 500ms。事实上显而易见，里面两层循环有极小的常数，所以 n^4 也是一个比较松的上界。

(C++)

#include <array>
#include <vector>
#include <iostream>
#include <cstdio>

typedef long long int ll;

const int maxn = 105;
const int maxm = 1000006;
const int p = 998244353;

int n, m, invm, invm2;
std::array<std::array<std::array<std::array<std::array<long long, 3>, maxn>, maxn>, 2>, maxn> f;
std::array<int, maxn> pos, sz;
std::array<long long, maxn> g;
std::array<std::vector<int>, maxn> e;

void dfs(const int u, const int fa);
void calc(const int u, const int fa);

int mpow(ll x, int y) {
  ll ret = 1;
  while (y) {
    if (y & 1) (ret *= x) %= p;
    y >>= 1;
    (x *= x) %= p;
  }
  return ret;
}

int main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(0);
  std::cin >> n >> m;
  invm = mpow(m, 998244351);
  invm2 = mpow(m - 2, 998244351);
  long long ans = m;
  for (int i = 1, u, v; i < n; ++i) {
    std::cin >> u >> v;
    e[u].push_back(v);
    e[v].push_back(u);
  }
  calc(1, 0);
  ll s1 = 0, s2 = 0;
  for (int i = n / 3 + 2; i <= n; ++i) {
    for (int j = 0; j <= n; ++j) {
      for (int k = 0; k <= 2; ++k) {
        (s1 += f[1][pos[1]][i][j][k]) %= p;
      }
    }
    for (int j = n / 3 + 2; j <= n; ++j) {
      for (int k = 0; k <= 2; ++k) {
        (s2 += f[1][pos[1]][i][j][k]) %= p;
      }
    }
  }
  ans = g[1];
  ans -= m * s1 % p;
  ans += (m - 1ll) * m / 2 % p * s2 % p;
  (ans += p) %= p;
  std::cout << ans << '\n';
}

void calc(const int u, const int fa) {
  f[u][pos[u]][1][0][0] = f[u][pos[u]][0][1][1] = 1; f[u][pos[u]][0][0][2] = m - 2;
  sz[u] = 1;
  g[u] = m;
  for (auto v : e[u]) if (v != fa) {
    calc(v, u);
    pos[u] ^= 1;
    (g[u] *= (m - 1ll) * invm % p * g[v] % p) %= p;
    for (int i = 0; i <= sz[u] + sz[v]; ++i) {
      for (int j = 0; j <= sz[u] + sz[v]; ++j) {
        f[u][pos[u]][i][j].fill(0);
      }
    }
    for (int i = sz[u]; i >= 0; --i) {
      for (int x = 0; x <= sz[v]; ++x) {
        for (int j = sz[u]; j >= 0; --j) {
          for (int y = 0; y <= sz[v]; ++y) {
            (f[u][pos[u]][i + x][j + y][0] += (f[u][pos[u] ^ 1][i][j][0]) * (f[v][pos[v]][x][y][1] + f[v][pos[v]][x][y][2]) % p) %= p;
            (f[u][pos[u]][i + x][j + y][1] += (f[u][pos[u] ^ 1][i][j][1]) * (f[v][pos[v]][x][y][0] + f[v][pos[v]][x][y][2]) % p) %= p;
            (f[u][pos[u]][i + x][j + y][2] += (f[u][pos[u] ^ 1][i][j][2]) * ((f[v][pos[v]][x][y][0] + f[v][pos[v]][x][y][1]) % p 
                + f[v][pos[v]][x][y][2] * invm2 % p * (m - 3) % p) % p) %= p;
          }
        }
      }
    }
    sz[u] += sz[v];
  }
}