题解：P17230 [Math×Girl²] 染色³
MadSamurai
·
·
题解
一种按常数行、常数列分类的解法。
下文把当前要求输出的网格写成 kA\times kB\times kC，其中 A=a 固定，1\le B\le b,1\le C\le c。
先把大网格按对齐的 k\times k\times k 小块切开。对齐的小块当然也是题目所说的子网格，所以每块恰好有一个黑格。把小块坐标记为 (x,y,z)，黑格在块内三个方向的偏移记成
(X_{y,z},Y_{x,z},Z_{x,y}).
为什么下标会少一个？沿某个方向把 k\times k\times k 窗口移动一格，移出去和新进来的两个截面必须含有相同数量的黑格。连续比较相邻窗口，就能得到第一维偏移不随 x 改变，另外两维同理。
于是一个方案其实就是三个矩阵
X\in[k]^{B\times C},\qquad
Y\in[k]^{A\times C},\qquad
Z\in[k]^{A\times B}.
两个粗网格点在若干维上不同的时候，合法条件等价于：至少有一个发生变化的维度，其对应偏移相同。直接拿这个条件做三张矩阵的容斥还是很难看，所以先固定 X。
设 X 恰好有 r 个常数行、s 个常数列。对每个 x，这一层还要选择一个长度为 C 的向量 Y_x 和一个长度为 B 的向量 Z_x。层内合法当且仅当 Y_x 是常向量或者 Z_x 是常向量，因此一层一共有
S_{B,C}=k^{B+1}+k^{C+1}-k^2
种状态，最后一项减掉两边都为常向量时的重复部分。
再检查两个 x 层之间的限制，会发现“能够放在同一个方案里”是一个等价关系。当 0\le r<B,0\le s<C 时，等价类只有下面三种：
这里 r 个常数行允许对应的 Z 分量自由变化，s 个常数列则允许对应的 Y 分量自由变化；两边都能变化的状态在常数 Y,Z 处交了一次，所以类大小里会出现 -1。
所有 A 层必须落在同一个等价类中，所以固定 X 后的方案数为
W^0_{B,C}(r,s)=k^2(k^r+k^s-1)^A
+(k^{B-r+1}-k^2)k^{rA}
+(k^{C-s+1}-k^2)k^{sA}.
边界稍微特殊一点：
W_{B,C}(B,C)=S_{B,C}^A,
W_{B,C}(B,0)=k^{BA+1}+k^{C+1}-k^2,\qquad
W_{B,C}(0,C)=k^{CA+1}+k^{B+1}-k^2.
除这三种外，不会出现 r=B 或 s=C。上面的 W^0 在后两个边界仍然正确，只在 (B,C) 需要补
k\left(S_{B,C}^A-W^0_{B,C}(B,C)\right),
因为同时全部为常数行、常数列的 X 一共有 k 张。
现在要数恰好有 r 个常数行、s 个常数列的 X。先指定 x 行和 y 列必须是常数，方案数记为 G_{B,C}(x,y)。这个数可以直接写：
G_{B,C}(x,y)=
\begin{cases}
k^{BC},&x=y=0,\\
k^{x+(B-x)C},&x>0,y=0,\\
k^{y+B(C-y)},&x=0,y>0,\\
k^{1+(B-x)(C-y)},&x>0,y>0.
\end{cases}
最后一种里，被指定的行列会在交点相遇，所以这些常数行列只能共用一种颜色。
二项式容斥以后，恰好数量为
N_{B,C}(r,s)=\binom Br\binom Cs
\sum_{u=0}^{B-r}\sum_{v=0}^{C-s}
(-1)^{u+v}\binom{B-r}{u}\binom{C-s}{v}
G_{B,C}(r+u,s+v).
于是直接做的话，答案就是
\operatorname{Ans}(B,C)=
\sum_{r=0}^{B}\sum_{s=0}^{C}N_{B,C}(r,s)W_{B,C}(r,s).
到这里小数据已经做完了，接下来只是想办法把这一大坨容斥一次算完。
对任意权值 W(r,s) 定义二维二项差分
D_W(i,j)=\sum_{r=0}^i\sum_{s=0}^j
(-1)^{i-r+j-s}\binom ir\binom jsW(r,s).
把前面 N 的式子代入并交换求和顺序，可以得到
\sum_{r,s}N_{B,C}(r,s)W(r,s)
=\sum_{i,j}\binom Bi\binom CjG_{B,C}(i,j)D_W(i,j).
这里把 G 换一种写法会更容易看出卷积：
G_{B,C}(i,j)=k^{(B-i)(C-j)}
\begin{cases}
1,&i=j=0,\\
k^i,&i>0,j=0,\\
k^j,&i=0,j>0,\\
k,&i>0,j>0.
\end{cases}
真正二维耦合的只有
F(r,s)=(k^r+k^s-1)^A.
其余部分只依赖 r 或只依赖 s，最后用二项式定理单独补掉即可。给数组除上阶乘：
X_{r,s}=\frac{F(r,s)}{r!s!},\qquad
E_{r,s}=\frac{(-1)^{r+s}}{r!s!},\qquad
K_{r,s}=\frac{k^{rs}}{r!s!}.
第一遍卷积 X*E 正好完成二维二项差分，第二遍再卷 K，对应 G 中剩余的 k^{(B-i)(C-j)}。记
V_{B,C}=B!C.
二维 NTT 不需要重新写，取步长 L=2c+1，把下标 (i,j) 压成 iL+j 后做普通一维卷积即可。需要注意不能把三个数组一次性放在频域里直接相乘：三个列次数相加可能超过 L，串到下一行。代码先算 X*E，裁掉 i>b 或 j>c 的部分，再与 K 卷一次，这样每次列次数都不超过 2c。
最后处理两条坐标轴、只依赖单个变量的部分和常数矩阵修正。令
t=k^A-1,\qquad p=k^B,\qquad q=k^C,
整理后每个输出位置是
\begin{aligned}
\operatorname{Ans}(B,C)={}&k^3\left(V_{B,C}-(q+t)^B-(p+t)^C-(p+q-1)^A\right)\\
&+k^{B+1}(q+t-k+1)^B+k^{C+1}(p+t-k+1)^C\\
&+k^2(k-1)k^{BC}+k(kp+kq-k^2)^A\\
&+(k^3-k^2)(k^{AB}+k^{AC}).
\end{aligned}
这条式子确实比较长，不过代码里没有更多隐藏容斥。一般情况总共做两次卷积，也就是六次 NTT，时间复杂度为
O(bc\log(bc)),
空间复杂度为 O(bc)。
实际写的时候不能真的对每个 (B,C) 再做六次快速幂。初始化 X 时把 (k^B+k^C-1)^A 留下来，另一项又有
(k^{B+1}+k^{C+1}-k^2)^A
=k^{2A}(k^{B-1}+k^{C-1}-1)^A,
所以也能复用左上角的值。其余四个幂分别沿行、列乘一次底数递推即可。另外 A=1 时直接使用上面的检查式；B=1 或 C=1 时答案可以化成
k^{ABC+1}(k^A+k^{BC}-k),
这几种情况都不需要进入 NTT。
还有一个常数问题：把二维下标直接压成一维时，即使 bc=10^6，窄长数据的 NTT 长度仍可能到 2^{23}。令 d=\min(b,c),e=\max(b,c)，当 d\le40 时，代码改为只在长的一维做长度约 2e 的 NTT；对每个频率，短的一维直接做 O(d^2) 卷积。两次二维卷积的这部分复杂度为
O(de\log e+d^2e),
在短边很小时比六遍巨大的一维 NTT 快很多。40 只是根据实现常数测出来的分界，附近的值不影响正确性。
附 AC 评测记录和代码