【集训队作业2018】复读机 题解

题意

求有多少个长为 $n$，值域为 $[1,k]$ 的序列满足 $\sum\limits_{i=1}^k[d\mid cnt_i]=k$，其中 $cnt_i$ 表示 $i$ 在序列中出现的次数，答案对 $19491001$ 取模。

$n\leq 10^9,k\leq 5\times 10^5,d\in\{1,2,3\}$，若 $d=3$ 则保证 $k\leq 10^3$。

$d=1$

大家都会！答案是 $k^n$。

$d=2$

写出每一种颜色的 $\mathrm{EGF}$：$F(x)=\sum\limits_{i\geq0}[d\mid i]\frac{x^i}{i!}$，那么 $ans=n![x^n]F^k(x)$。

这个 $[d\mid i]$ 看起来很不爽，让我们找不到封闭形式，直接把他单位根反演掉：

$F(x)=\sum\limits_{i\geq0}\frac1d\sum\limits_{j=0}^{d-1}\omega_d^{ij}\frac{x^i}{i!}=\frac1d\sum\limits_{i=0}^{d-1}\sum\limits_{j\geq0}\frac{(\omega_d^ix)^j}{j!}=\frac1d\sum\limits_{i=0}^{d-1}e^{\omega_d^ix}$

这样就有封闭形式了，我们直接把 $d=2$ 的时候的单位根代入：

$ans=n^k=\frac{n!}{d^k}[x^n]\sum\limits_{i=0}^k\dbinom kie^{ix}e^{-(k-i)x}$

$=\frac{n!}{d^k}[x^n]\sum\limits_{i=0}^k\dbinom kie^{(2i-k)x}=\frac{n!}{d^k}[x^n]\sum\limits_{i=0}^k\dbinom ki\sum\limits_{j\geq0}\frac{((2i-k)x)^j}{j!}$

$=\frac{n!}{d^k}\sum\limits_{i=0}^k\dbinom ki\frac{(2i-k)^n}{n!}=\frac{1}{d^k}\sum\limits_{i=0}^k\dbinom ki(2i-k)^n$

枚举 $i$，时间复杂度 $O(k\log n)$。

$d=3$

我们延续 $d=2$ 的思路，把 $d=3$ 时的单位根代入：

$ans=n^k=\frac{n!}{d^k}[x^n]\sum\limits_{i=0}^k\dbinom ki\sum\limits_{j=0}^{k-i}\dbinom{k-i}je^{ix}e^{j\omega_3^1x}e^{(k-i-j)\omega_3^2x}$

$=\frac{n!}{d^k}[x^n]\sum\limits_{i=0}^k\dbinom ki\sum\limits_{j=0}^{k-i}\dbinom{k-i}je^{(i+j\omega_3^1+(k-i-j)\omega_3^2)x}=\frac{n!}{d^k}[x^n]\sum\limits_{i=0}^k\dbinom ki\sum\limits_{j=0}^{k-i}\dbinom{k-i}j\sum\limits_{p\geq0}\frac{((i+j\omega_3^1+(k-i-j)\omega_3^2)x)^p}{p!}$

$=\frac{n!}{d^k}\sum\limits_{i=0}^k\dbinom ki\sum\limits_{j=0}^{k-i}\dbinom{k-i}j\frac{(i+j\omega_3^1+(k-i-j)\omega_3^2)^n}{n!}=\frac{1}{d^k}\sum\limits_{i=0}^k\dbinom ki\sum\limits_{j=0}^{k-i}\dbinom{k-i}j(i+j\omega_3^1+(k-i-j)\omega_3^2)^n$

枚举 $i,j$，时间复杂度 $O(k^2\log n)$。

代码

冷知识：$19491001$ 的原根是 $7$。

#include<iostream>
#include<cstdio>
using namespace std;
const int mod=19491001,g=7;
int n,k,d,fac[500001],inv[500001],w[3],ans;
inline int pw(int a,int b)
{
    int res=1;
    while(b)
    {
        if(b&1)
            res=1ll*res*a%mod;
        b>>=1;
        a=1ll*a*a%mod;
    }
    return res;
}
inline int c(int x,int y)
{
    if(x<y||x<0||x<0)
        return 0;
    return 1ll*fac[x]*inv[y]%mod*inv[x-y]%mod;
}
int main()
{
    cin>>n>>k>>d;
    if(d==1)
    {
        cout<<pw(k,n)<<'\n';
        return 0;
    }
    fac[0]=inv[0]=1;
    for(int i=1;i<=k;++i)
        fac[i]=1ll*fac[i-1]*i%mod;
    inv[k]=pw(fac[k],mod-2);
    for(int i=k-1;i>=1;--i)
        inv[i]=1ll*inv[i+1]*(i+1)%mod;
    if(d==2)
    {
        for(int i=0;i<=k;++i)
            ans=(ans+1ll*c(k,i)*pw((2*i-k+mod)%mod,n)%mod)%mod;
        cout<<1ll*ans*pw(pw(2,k),mod-2)%mod<<'\n';
        return 0;
    }
    w[0]=1;
    w[1]=pw(g,(mod-1)/3);
    w[2]=1ll*w[1]*w[1]%mod;
    for(int i=0;i<=k;++i)
    {
        int res=0;
        for(int j=0;j<=k-i;++j)
            res=(res+1ll*inv[j]*inv[k-i-j]%mod*pw((i+1ll*w[1]*j%mod+1ll*w[2]*(k-i-j)%mod)%mod,n)%mod)%mod;
        ans=(ans+1ll*inv[i]*res%mod)%mod;
    }
    cout<<1ll*fac[k]*pw(pw(3,k),mod-2)%mod*ans%mod<<'\n';
    return 0;
}