P3784 [SDOI2017] 遗忘的集合

传送门

description

对于一个元素都 \(\leq n\) 的正整数集合 \(S\)（不含相同元素），\(f(i)\) 表示使用集合 \(S\) 里的数加和为 \(i\) 的方案数，每个元素可以被使用多次，两个方案不同当且仅当存在一个元素在两种方案中使用次数不同。

现给定 \(n\) 和 \(f(i),1\leq i\leq n\)。

求出集合 \(|S|\) 的大小并给出字典序最小的 \(S\)。

solution

先反着考虑给定集合 \(S\) 怎么求出 \(f(i)\)。

设 \(f_{i,j}\) 表示用 \(S\) 里前 \(i\) 个元素构成和为 \(j\) 的方案数，设 \(S\) 中第 \(i\) 个元素为 \(a_i\)。有转移：

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

• \(f_{i,j}=\sum\limits_{k\ge 0} f_{i-1,j-ka_i}\)

设 \(F_i(x)\) 是 \(\{f_{i,j}\}_{j=0}^{+ \infty}\) 的生成函数，根据转移式，有 \(F_{|S|}(x)=\prod\limits_{i=1}^{|S|}(1+x^{a_i}+x^{2a_i}+x^{3a_i}+\dots)\)。

根据经典套路，右边取 \(\ln\) 得，\(\ln(F_{|S|}(x))=\sum\limits_{i=1}^{|S|}\ln(1+x^{a_i}+x^{2a_i}+x^{3a_i}+\dots)=\sum\limits_{i=1}^{|S|}\ln(\dfrac{1}{1-x^{a_i}})\)。

根据 \(\ln(1-x)=-\sum\limits_{i\ge 1}\dfrac{x^i}{i}\) 得，\(\ln(F_{|S|}(x))=\sum\limits_{i=1}^{|S|}\sum\limits_{j\ge 1} \dfrac{x^{ja_i}}{j}\)。

右边调换一下求和顺序，\(\sum\limits_{p=1}x^p\sum\limits_{d\mid p} A_{d}\dfrac{d}{p}\)，其中 \(A_d\) 表示 \(d\) 是否在集合 \(S\) 中。

归纳地，假设已经求出所有 \(A_{i},1\leq i<t\)，我们当然能求出 \(A_t\)，直接枚举从小到大 \(d\) 并枚举它的倍数即可，时间复杂度 \(O(n\ln n)\)。

带上多项式求 \(\ln\)，时间复杂度 \(O(n\log n)\)。

我用的之前写的 3 模数 NTT 求逆板子，常数比较大，不过轻松跑，一发就过了。（这题 5s 时限）。

顺便说一下，根据我们构造 \(S\) 的方法可以知道，如果数据保证有解，那么 \(S\) 一定唯一。

code

#include<bits/stdc++.h>

using namespace std;

typedef long long E;
typedef unsigned __int128 lint;
const E mod1=998244353,mod2=469762049,mod3=1004535809,_g=3;
const int N=(1<<19)+9;
const lint P=(lint)mod1*mod2*mod3;

struct Ei{
    E a,b,c;

    Ei(){

    }

    Ei (E x,E y,E z){
        a=x,b=y,c=z;
    }

    Ei(E x){
        a=b=c=x;
    }

    friend Ei operator +(const Ei &x,const Ei &y){
        Ei z;
        z.a=x.a+y.a>=mod1?x.a+y.a-mod1:x.a+y.a;
        z.b=x.b+y.b>=mod2?x.b+y.b-mod2:x.b+y.b;
        z.c=x.c+y.c>=mod3?x.c+y.c-mod3:x.c+y.c;
        return z;
    }

    friend Ei operator -(const Ei &x,const Ei &y){
        Ei z;
        z.a=x.a<y.a?x.a-y.a+mod1:x.a-y.a;
        z.b=x.b<y.b?x.b-y.b+mod2:x.b-y.b;
        z.c=x.c<y.c?x.c-y.c+mod3:x.c-y.c;
        return z;
    }

    friend Ei operator *(const Ei &x,const Ei &y){
        Ei z;
        z.a=x.a*y.a%mod1;
        z.b=x.b*y.b%mod2;
        z.c=x.c*y.c%mod3;
        return z;
    }
};

inline E ksm(E a,E b,E mod){
    E ret=1;
    if(a>=mod) a%=mod;
    if(b>=mod-1||b<0) b=(b%(mod-1)+mod-1)%(mod-1);
    while(b){
        if(b&1) ret=ret*a%mod;
        a=a*a%mod;
        b>>=1;
    }
    return ret;
}


const lint inv1=ksm(mod2*mod3%mod1,mod1-2,mod1);
const lint inv2=ksm(mod1*mod3%mod2,mod2-2,mod2);
const lint inv3=ksm(mod1*mod2%mod3,mod3-2,mod3);
const lint p1=P/mod1*inv1%P,p2=P/mod2*inv2%P,p3=P/mod3*inv3%P;
E mod=1e9+7;

void wr(lint x){
    if(x>9) wr(x/10);
    putchar('0'+x%10);
}

void CRT(Ei *a,int n){
    for(int i=0; i<n; i++){
        lint ans=0;
        ans=(ans+(lint)a[i].a*p1);
        ans=(ans+(lint)a[i].b*p2);
        ans=(ans+(lint)a[i].c*p3)%P;
        a[i]=ans%mod;
    }
}

int rev[N],lstn;

inline void prework(int n){
    if(lstn==n) return ;
    for(int i=1; i<n; i++) rev[i]=(rev[i^(i&-i)]|(1<<(__lg(n)-1-__lg(i&-i))));
    return lstn=n,void();
}

void ntt(Ei *p,int n,int op){
    prework(n);
    for(int i=1; i<n; i++){
        if(i<rev[i]) swap(p[i],p[rev[i]]);
    }

    for(int i=2; i<=n; i<<=1){
        int len=i>>1;
        Ei w=(Ei){ksm(_g,op*(mod1-1)/i,mod1),ksm(_g,op*(mod2-1)/i,mod2),ksm(_g,op*(mod3-1)/i,mod3)};
        for(int pos=0; pos<n; pos+=i){
            Ei now=1ll;
            for(int j=pos; j<pos+len; j++,now=now*w){
                Ei u=p[j],v=now*p[j+len];
                p[j]=u+v,p[j+len]=u-v;
            }
        }
    }

    if(op==-1){
        Ei tmp=(Ei){ksm(n,mod1-2,mod1),ksm(n,mod2-2,mod2),ksm(n,mod3-2,mod3)};
        for(int i=0; i<n; i++){
            p[i]=p[i]*tmp;
        }
    }
}

void px(Ei *a,Ei *b,int n){
    for(int i=0; i<n; i++) a[i]=a[i]*b[i];
}

#define clr(f,n) memset(f,0,sizeof(Ei)*n)
#define cpy(f,g,n) memcpy(f,g,sizeof(Ei)*n)

void polyinv(Ei *a,Ei *b,int n){
    static Ei now[N],r[N];
    b[0]=ksm(a[0].a,mod-2,mod);
    for(int len=2; len<=n; len<<=1){
        int t=len>>1;
        for(int i=0; i<t; i++){
            r[i]=b[i].a*2%mod;
        }
        cpy(now,a,len);
        ntt(now,len<<1,1),ntt(b,len<<1,1);
        px(now,b,len<<1);
        ntt(now,len<<1,-1);
        clr(now+len,len);
        CRT(now,len);
        for(int i=0; i<(len); i++) now[i]=(mod-now[i].a)%mod;

        ntt(now,len<<1,1);
        px(now,b,len<<1);
        ntt(now,len<<1,-1);
        clr(now+len,len);
        CRT(now,len);
        for(int i=0; i<len; i++){
            b[i]=(r[i].a+now[i].a>=mod?r[i].a+now[i].a-mod:r[i].a+now[i].a);
        }
        clr(b+len,len);
    }
    clr(now,n<<1),clr(r,n<<1);
}


inline void polyln(Ei *a,Ei *b,int n){
  static Ei tmp[N];
  clr(b,n<<1); clr(tmp,n<<1);
  polyinv(a,tmp,n);
  for(int i=0; i<n; i++){
    b[i]=a[i+1].a*(i+1)%mod;
  }
  ntt(b,n<<1,1),ntt(tmp,n<<1,1);
  px(b,tmp,n<<1);
  ntt(b,n<<1,-1);
  clr(b+n,n);
  CRT(b,n);
  for(int i=n-1; i; i--){
    b[i]=b[i-1].a*ksm(i,mod-2,mod)%mod;
  }
  b[0]=0;
}

E a[N],b[N],n,m,INV[N];
Ei A[N],B[N];

int main(){

  cin>>n>>mod;
  INV[1]=1;
  for(int i=2; i<=n; i++) INV[i]=(mod-mod/i)*1ll*INV[mod%i]%mod;
  for(int i=1; i<=n; i++) scanf("%lld",a+i);
  a[0]=1;
  for(int i=0; i<=n; i++) A[i]=a[i];
  int t;
  for(t=1;t<=n;t<<=1) ;
  polyln(A,B,t);
  for(int i=1; i<=n; i++) b[i]=B[i].a;

  for(int i=1; i<=n; i++){
    for(int j=i*2; j<=n; j+=i){
      b[j]=(b[j]-b[i]*INV[j/i]%mod+mod)%mod;
    }
  }

  vector<int> ans;
  for(int i=1; i<=n; i++){
    if(b[i]!=0) ans.emplace_back(i);
  }

  cout<<ans.size()<<endl;
  for(auto p:ans) printf("%lld ",p);

  return 0;
}

P.S.

好像代码里递推部分最开始写成 \(\log^2\) 了（调和级数套快速幂），改成线性预处理逆元之后快了 10s 多。