题解：AT_abc387_f [ABC387F] Count Arrays
shiminghai
·
·
题解
事情的起因是我没看见 $M \le 2025$ 这个限制，因此本文讨论的是 $M \le 10^9$ 的做法，官方做法请看其它题解。
### 题目分析：
- 首先建树是肯定需要的，让每一个 $A_i$ 向 $i$ 建一条有向边，不难发现这是一个基环树森林，问题也就变为了求所有叶子节点值小于等于父节点的方案数。
- 接下来考虑单独一个基环树，结果等于每个树的方案数之积（不同基环树之间无相互影响）。
- 令 $x_{u,t}$ 表示 $u$ 节点值为 $t$ 时当前子树方案数。转移比较简单，$x_{u,t}=\prod_{v\in son(u)}\left(\sum_{s=1}^{t} x_{v,s}\right)$，就是每一个子树的方案数之积。
- 考虑用 NTT 优化，令 $dp_{u,i}$ 表示 $u$ 节点多项式的 $i$ 次项系数，其中的多项式是 $x_{u,t}$ 与方案数的表达式，也就是将节点取值带入多项式计算能得出方案数。
- 如果你没看懂，举个例子：一个叶子 $u$，它的值无论是多少，方案数都是 $1$，因此多项式就是 $t=0t^{2025}+...+0t+1$，因此只有 $dp_{u,0}=1$，也就是 $0$ 次项系数为 $1$，其它均为 $0$。
- 这样，节点 $u$ 的多项式就是每个儿子节点的多项式的前缀和的乘积（前缀和对应原转移的 $\sum_{s=1}^{t} dp_{v,s}$，乘积对应 $\prod_{v\in son(u)}$）。
- 至此我们通过 $O(n^2\log n)$ 的做法通过了此题，复杂度优化了 `棍母`，码量 $+=\infty$。
- ::::info[到此为止？]
这个做法看着像是是有前途的，我肯定优化不了，因此我去问了 gpt，它说能将复杂度降到小于 $n^2$。
:::info[GPT 的解释：]
可以做到严格次二次。对任意树形，重链分块、局部泰勒展开和快速矩阵乘法可以做到 $(O_\varepsilon(N^{13/8+\varepsilon}))$ 次模域运算，复杂度不随 $M$ 增长。指数用到了任意域上 $\omega\le 9/4$ 的矩阵乘法结果。（[arxiv.org](https://arxiv.org/pdf/2611.09716#1)）\
关键是一整段重链只做一次大多项式转移。段内的转移先合成低阶微分算子，再用同一张小矩阵批量处理输入多项式的局部展开。
### 换基底，把前缀和变成移位
原题做法：设 $F_{u,t}$ 是 $u$ 的子树中，根的取值至多为 $t$ 的方案数，则
$$
F_u=S\left(\prod_{v\in son(u)}F_v\right),
S_{f_t}=\sum_{i=1}^{t}f(i).
$$
最后求各根的 $F_{r,m}$ 并相乘。
用上升二项式基底 $b_{k,t}=\binom{t+k-1}{k}$，把 $\sum_k c_k b_{k,t}$ 表示成 $g(z)=\sum_k c_k z^k$。因为
$$
S_{b_k}=b_{k+1},t b_k=(k+1)b_{k+1}-k b_k,
$$
所以前缀和对应乘 $z$，乘 $t$ 对应算子
$$
D=z+(z^2-z)\frac{d}{dz}.
$$
纯链因此可以直接给出 $\binom{M+\ell-1}{\ell}$。而一般重链上，记 $Q_{u,t}$ 为所有轻儿子的 $F_{v,t}$ 的乘积，那么节点 $u$ 的转移就是 $g\mapsto z Q_{u,D} g$。
### 矩阵处理整段重链
给节点赋权 $w_u=1+\sum_{v \in{son}_{\text{light}}(u)}|{size}(v)|$。一段总权不超过 $B$ 的重链，合成算子 $L$ 的微分阶数至多 $B$，且它至多把多项式次数提高 $B$。\
对任意采样点 $a$，取
$$
g_a=g\bmod(z-a)^{2B}.
$$
每求一次导数，零点重数至多下降 $1$，因此
$$
Lg\equiv Lg_a\pmod{(z-a)^B}.
$$
这里最有用的是，每个 $g_a$ 都以同一个全局变量 $z$ 的幂基底表示，次数小于 $2B$。于是所有 $Lg_a$ 都能用同一张 $O(B)\times O(B)$ 的矩阵乘出来。若整条重链链头的子树大小为 $s$，取 $O(s/B)$ 个采样点，最后用多项式 CRT 恢复 $Lg$。取余和恢复用乘积树，合计 $\widetilde O(s)$。（[mathexp.eu](https://mathexp.eu/bostan/publications/BeBoHo12.pdf)）\
这张矩阵也能分治构造。对总权为 $w$ 的算子，有
$$
L(z^k)=\sum_{j=0}^{w}p_j(k)z^{k+j},\deg p_j\le w.
$$
这是因为 $D$ 作用在 $z^k$ 上只产生两个相邻次数，系数都是 $k$ 的一次式。于是矩阵的每条对角线都能从 $O(w)$ 个连续样本插值，并延长到合并所需的列数。合并两段时，把延长后的作用矩阵相乘即可，整块构造耗时 $\widetilde O(B^\omega)$。而单节点的列可以通过在 $t=0,-1,-2...$ 上做二项式变换、点乘 $Q_u$、再逆变换来生成（因为 $b_k(-i)=(-1)^k\binom{i}{k}$）。
### 证明复杂度：
一条重链上的权重之和恰好等于链头的子树大小 $s$。取 $B=\lceil\sqrt s\rceil$，按总权贪心分块，得到 $O(\sqrt s)$ 块。单个权重超过 $B$ 的节点直接做一次多项式乘法和前缀和。其余每块的采样点数和矩阵维度都是 $O(\sqrt s)$，所以整条链耗时
$$
\widetilde O\left(s^{(\omega+1)/2}\right).
$$
每个点只会被计入 $O(\log N)$ 个重链链头的子树，因为跨一条轻边，子树大小至少减半。对所有重链求和即得 $\widetilde O(N^{(\omega+1)/2})$。最后用 $b_k(M)$ 线性求值即可。
:::
这样最劣情况是一颗梳子树，即一条长主链，沿途每个节点各挂一个叶子，复杂度会降到 $O(n^{1.904})$（我不知道为啥反正是 gpt 说的 qwq）。\
代码我让它写了一版，应该没问题。\
[GPT 代码](https://atcoder.jp/contests/abc387/submissions/80743012)为[目前最优解](https://atcoder.jp/contests/abc387/submissions?f.LanguageName=\&f.Status=AC\&f.Task=abc387_f\&f.User=\&orderBy=time_consumption)。
::::
### 代码实现：
- 用一个能处理前缀和与乘积的 NTT。
- 以环的每一个点为根做一遍。
- 标准 Bernoulli 约定 $B_1=-\frac12$，但前缀和公式 $\sum_{s=1}^t s^k=\frac{1}{k+1}\sum_{j=0}^{k}\binom{k+1}{j}B_j^+,t^{k+1-j}$ 需要 $B_1^+=+\frac12$（$B_1$ 的符号为 $+$）。因此代码在计算完后对 $B_1^+$ 取负了。
- 其它细节详见代码与注释。
:::info[code：]
没想到吧，我居然写了代码。
```cpp
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int mod=998244353,G=3,N=2050,M=1<<13;
int n,m,A[N],du[N];//du出度;
int dp[N][N],dplen[N];//dp[i][j]第i个点多项式的j次项系数,dplen多项式长度
int ntt_a[M],ntt_b[M];//NTT
int buf1[N],buf2[N],buf3[N];//暂时存储当前累计多项式的数组
int jc[N+5],njc[N+5],bern[N+5];//jc阶乘,njc阶乘逆元,bern存Bernoulli数
int c_buf[N],e_buf[N],cr_buf[N],conv_buf[2*N];//Bernoulli部分
/*
c_buf[i]=多项式i次项系数×i!
e_buf[i]=bern[i]/(i!)
cr_buf[i] 下标翻转的c_buf数组,卷积用
conv_buf[i]=cr_buf[i]*e_buf[i]的卷积结果
*/
int order[N],cyc[N];//order遍历顺序,cyc存环上元素下标
vector<int>g[N];
bool vis[N],h[N];//vis是否访问过,h是否在环上
int qpow(int x,int y)
{
int res=1;
while(y)
{
if(y&1)
res=res*x%mod;
x=x*x%mod;
y>>=1;
}
return res;
}
void ntt(int*a,int n,bool flag)//flag 正/逆DFT
{
for(int i=1,j=0;i<n;i++)
{
int bit=n>>1;
for(;bit&j;bit>>=1)
j^=bit;
j^=bit;
if(i<j)
swap(a[i],a[j]);
}
for(int len=2;len<=n;len<<=1)
{
int wlen=qpow(G,(mod-1)/len),half=len>>1;;//计算本原len次单位根,half基2分治
if(flag)
wlen=qpow(wlen,mod-2);//转逆元
for(int i=0;i<n;i+=len)
{
int w=1;//旋转因子
for(int k=0;k<half;k++)
{
int u=a[i+k],v=a[i+k+half]*w%mod;
a[i+k]=u+v;
if(a[i+k]>=mod)
a[i+k]-=mod;
a[i+k+half]=u-v;
if(a[i+k+half]<0)
a[i+k+half]+=mod;
w=w*wlen%mod;
}
}
}
if(flag)//逆DFT需要的/n
{
int inv_n=qpow(n,mod-2);
for(int i=0;i<n;i++)
a[i]=a[i]*inv_n%mod;
}
}
int poly_mul(const int*a,int la,const int* b,int lb,int*out)//a,b多项式相乘
{
if(la==0||lb==0)
return 0;
int rl=la+lb-1,n=1;
while(n<rl)//把n翻倍到不小于rl的最小2的幂,因为NTT蝶形运算要求数组长度是2的幂。
n<<=1;
for(int i=0;i<n;i++)//多余位置用0补位
{
ntt_a[i]=(i<la)?a[i]:0;
ntt_b[i]=(i<lb)?b[i]:0;
}
ntt(ntt_a,n,0);//正DFT
ntt(ntt_b,n,0);
for(int i=0;i<n;i++)
ntt_a[i]=ntt_a[i]*ntt_b[i]%mod;//系数直接相乘
ntt(ntt_a,n,1);//逆DFT
for(int i=0;i<rl;i++)
out[i]=ntt_a[i];
return rl;
}
int poly_eval(const int*a,int la,int t)//计算a多项式变量为t时的值
{
int v=0;
for(int i=la-1;i>=0;i--)
v=(v*t+a[i])%mod;
return v;
}
void ycl()//预处理阶乘
{
jc[0]=1;
for(int i=1;i<=N;i++)
jc[i]=jc[i-1]*i%mod;
njc[N]=qpow(jc[N],mod-2);
for(int i=N;i;i--)
njc[i-1]=njc[i]*i%mod;
bern[0]=1;
for(int i=1;i<N;i++)
{
int s=0;
for(int k=0;k<i;k++)
{
int c=jc[i+1]*njc[k]%mod*njc[i+1-k]%mod;
s=(s+c*bern[k])%mod;
}
bern[i]=(mod-s*qpow(i+1,mod-2)%mod)%mod;
}
bern[1]=(mod-bern[1])%mod;//取反
}
int poly_sum(const int* a,int la,int*out)//对多项式a求前缀和
{
int d=la-1;
for(int k=0;k<=d;k++)
c_buf[k]=a[k]*jc[k]%mod;
for(int r=0;r<=d;r++)//Bernoulli因子
e_buf[r]=bern[r]*njc[r]%mod;
for(int i=0;i<=d;i++)//反序cr
cr_buf[i]=c_buf[d-i];
poly_mul(cr_buf,d+1,e_buf,d+1,conv_buf);//大力NTT卷积
out[0]=0;
for(int m=1;m<=d+1;m++)
out[m]=conv_buf[d+1-m]*njc[m]%mod;//还原时除以原先乘的阶乘
return d+2;//前缀和把多项式次数加了1,系数个数=次数+1,因此+2
}
signed main()
{
ycl();
int ord_cnt=0,ans=1;//order的元素个数,ans总方案数
scanf("%lld%lld",&n,&m);
for(int i=0;i<n;i++)
{
scanf("%lld",&A[i]);
A[i]--;
du[A[i]]++;
g[A[i]].push_back(i);//值向下标建边
}
for(int i=0;i<n;i++)//枚举起点,拨叶子找环
{
if(du[i]!=0||vis[i])
continue;
int u=i;//u为叶子
while(!vis[u])
{
vis[u]=1;
order[ord_cnt++]=u;//按order处理每个点被处理时所有叶子结点一定被处理完了
int v=A[u];//v为父亲
du[v]--;
if(du[v]==0)
u=v;
else
break;
}
}
for(int i=0;i<n;i++)
h[i]=!vis[i];
//非环节点从叶到根DP,处理除环以外所有点
for(int i=0;i<ord_cnt;i++)
{
int u=order[i];
int cur_len = 1;//当前节点多项式系数个数(不含前导0)
buf1[0]=1;//至少有一个方案,0次项设为1
bool first=1;//是否为第一个儿子
for(int v:g[u])
{
int poly_len=poly_sum(dp[v],dplen[v],buf2);//对buf2做前缀和并返回系数长度
if(first)//第一个儿子就做前缀和再拷贝到buf1即可,其它儿子需要合并
{
for(int i=0;i<poly_len;i++)
buf1[i]=buf2[i];
cur_len=poly_len;
first=0;
}
else
{
cur_len=poly_mul(buf1,cur_len,buf2,poly_len,buf3);//更改buf3为两多项式乘积,同时返回长度
for(int i=0;i<cur_len;i++)//不直接传buf1因为同时有修改和查询,加一个中间数组存结果
buf1[i]=buf3[i];
}
}
for(int i=0;i<cur_len;i++)
dp[u][i]=buf1[i];
dplen[u]=cur_len;
}
//处理环并统计答案
for(int i=0;i<n;i++)
vis[i]=0;
for(int i=0;i<n;i++)
{
if(!h[i]||vis[i])
continue;
int cyc_len=0;//当前环上元素个数
int u=i;
cyc[cyc_len++]=i;
vis[i]=1;
u=A[i];
while(u!=i)
{
cyc[cyc_len++]=u;
vis[u]=1;
u=A[u];
}
//找到一个完整环了,处理环上点的多项式
for(int id=0;id<cyc_len;id++)
{
int u=cyc[id],c_len=1;//当前多项式系数个数(同样不含前导0)
buf1[0]=1;
bool first=1;//下面多项式合并和上面的差不多,就不注释了
for(int x:g[u])
{
if(h[x])//这里说一下,环上的儿子是不统计的
continue;
int poly_len=poly_sum(dp[x],dplen[x],buf2);
if(first)
{
for(int j=0;j<poly_len;j++)
buf1[j]=buf2[j];
c_len=poly_len;
first=0;
}
else
{
c_len=poly_mul(buf1,c_len,buf2,poly_len,buf3);
for(int j=0;j<c_len;j++)
buf1[j]=buf3[j];
}
}
for(int j=0;j<c_len;j++)
dp[u][j]=buf1[j];
dplen[u]=c_len;
}
//将环上所有点合并到一起,因为取值相同所以不用sum的前缀和
int poly_len=1;//当前合并时多项式系数个数
buf1[0]=1;
bool first=1;
for(int id=0;id<cyc_len;id++)
{
int u=cyc[id];
if(first)
{
for(int j=0;j<dplen[u];j++)
buf1[j]=dp[u][j];
poly_len=dplen[u];
first=0;
}
else
{
poly_len=poly_mul(buf1,poly_len,dp[u],dplen[u],buf3);
for(int j=0;j<poly_len;j++)
buf1[j]=buf3[j];
}
}
int ans_len=poly_sum(buf1,poly_len,buf2);//对多项式做前缀和,随便存到一个buf里,接下来带入m值
ans=(ans*poly_eval(buf2,ans_len,m))%mod;//带入值为m求出方案数,相乘
}
printf("%lld\n",ans);//齐活!
return 0;
}
```
:::
### 复杂度：
每个节点在每次合并时前缀和导致次数 $+1$，相乘只是次数相加，最多 $n$ 次合并。因此每个节点最多贡献 $n$ 的次数，总次数就是在 $n^2$ 量级。
时间复杂度 $O(n^2\log n)$（还有 NTT 的复杂度）。\
空间复杂度 $O(n)$。\
[AC 记录](https://atcoder.jp/contests/abc387/submissions/80746018)。