CF140E New Year Garland

题目链接。

Description

有 \(m\) 种不同的小球，要装饰成一个 \(n\) 层的树，每层有 \(L_i\) 个小球，要求：

• 每一层内部相邻小球不相同。
• 相邻两层的小球颜色集合不同。

求方案数。

Solution

当你划分完毕每层的颜色集合后，每层的方案是独立的，可以相乘。

考虑先每一层单独处理：

状态设计

设 \(f_{i, j}\) 为前 \(i\) 个小球，用了 \(j\) 种颜色方案数。

初始状态

\(f_{1,1} = 1\)

状态转移

• 考虑最后一个球是靠新增的颜色来的，即 \(f_{i, j} \gets f_{i - 1, j - 1}\)
• 考虑最后一个小球是用的之前的颜色，但是只能填 \(j - 1\) 个颜色，因为不能和相邻左侧的颜色相同。

综上：\(f_{i, j} = f_{i - 1, j - 1} + (j - 1) * f_{i - 1, j}\)

时间复杂度

\(O(L^2)\) 可以预处理所有情况

把答案叠加起来

发现当且仅当两层的颜色集合的个数相等，且正好选中相同的才会重合，不妨考虑容斥原理。如果按层考虑，这一层还要依赖于上一层的选择，所以还得来个 DP。

状态设计

设 \(g_{i, j}\) 为前 \(i\) 层，第 \(i\) 层颜色集合为 \(j\) 的方案数。

状态转移

• 首先不考虑限制，枚举上面一层放的颜色数 \(k \le L_{i - 1}\)，随便放：\(g_{i, j} \gets \sum g_{i - 1, k} \times f_{L_i, j} \times A_{m}^{j}\)
• 考虑将两次颜色集合相同的情况减掉，\(g_{i, j} \gets -g_{i - 1, j} \times f_{L_i, j} \times A_{j}^{j}\)。

综上：\(g_{i, j} = f_{L_i, j} \times (\sum g_{i - 1, k} \times A_{m}^{j} - g_{i - 1, j}\times A_{j}^{j})\)

复杂度

那个 \(A\) 可以预处理之后 \(O(1)\) 回答。

看似复杂度爆炸，但是每一层的 \(j \le L[i]\)，而 \(\sum_{i=1}^{n} L[i] \le 10 ^ 7\)，所以总复杂度是 \(O(10^7)\) 的不会爆炸。

Tips

1. 注意 \(g\) 数组 \(i = 1\) 的时候要特判

2. \(p\) 不一定是质数，不一定有逆元。但是我们需要的 \(A\) 很特殊，要么是 \(A_{j}^{j} = j!\)；要么是 \(A_{m}^{j}\)，其实是 \(m \times (m - 1) \times ... \times (m - j + 1)\)，利用 \(A_{m}^{j} = A_{m}^{j - 1} \ \times (m - j + 1)\) 进行 \(O(m)\) 预处理即可。

3. \(g\) 要用滚动数组否则空间会爆炸。

4. \(g_{i - 1, j}\) 可能不存在，当 \(j > L_{i - 1}\) 时，由于滚动数组没有覆盖，可能会出错，需要特判。

做完这道题想了想，为什么不能直接求 \(g\) 的时候求排列意义下的方案呢？不能的原因是之后容斥需要精准的找到上下矛盾的情况，不仅仅是颜色数意义还得恰好选的一样。如果直接求就无法区分，或者要涉及除法，然而这题不保证 \(P\) 是质数不一定有逆元，所以要规避掉这种需要除法的情况，用乘法代替除法。

#include <iostream>
#include <cstdio>
 
using namespace std;
 
typedef long long LL;
 
const int N = 1000005, M = 5005;
 
int n, m, P, fact[M], A[M], l[N], L, f[M][M], g[2][M];
 
int main() {
	scanf("%d%d%d", &n, &m, &P);
	for (int i = 1; i <= n; i++) scanf("%d", l + i), L = max(L, l[i]);
	fact[0] = A[0] = 1;
	for (int i = 1; i <= L; i++) fact[i] = (LL)fact[i - 1] * i % P;
	for (int i = 1; i <= L; i++) A[i] = (LL)A[i - 1] * (m - i + 1) % P;
	f[1][1] = 1;
	for (int i = 2; i <= L; i++) {
		for (int j = 1; j <= i; j++) {
			f[i][j] = (f[i - 1][j - 1] + (LL)f[i - 1][j] * (j - 1)) % P;
		}
	}
	int s = 0;
	for (int j = 1; j <= min(m, l[1]); j++)
		g[1][j] = (LL)f[l[1]][j] * A[j] % P, (s += g[1][j]) %= P;
 
	for (int i = 2; i <= n; i++) {
		for (int j = 1; j <= min(m, l[i]); j++) {
			g[i & 1][j] = ((LL)s * A[j] % P - (j > l[i - 1] ? 0 : (LL)g[(i - 1) & 1][j] * fact[j] % P) + P) % P * f[l[i]][j] % P;
		}
		s = 0;
		for (int j = 1; j <= min(m, l[i]); j++) (s += g[i & 1][j]) %= P;
	}
 
	printf("%d\n", s);
	return 0;
}