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} imes f_{L_i, j} imes A_{m}^{j})
• 考虑将两次颜色集合相同的情况减掉，(g_{i, j} gets -g_{i - 1, j} imes f_{L_i, j} imes A_{j}^{j})。

综上：(g_{i, j} = f_{L_i, j} imes (sum g_{i - 1, k} imes A_{m}^{j} - g_{i - 1, j} imes 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 imes (m - 1) imes ... imes (m - j + 1))，利用 (A_{m}^{j} = A_{m}^{j - 1} imes (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
", s);
	return 0;
}