BZOJ 3771 Triple

思路

河神是真的只可能拿一把、两把、三把斧子
设\(A(x)\)为所有斧子代价的生成函数
只拿一把斧子的方案数为\(A(x)\)
拿了两把斧子的方案数不是\(A^2(x)\)，因为可能会有拿重复的，手动容斥一下，设\(B(x)\)为所有斧子代价的两倍的生成函数（相当于取了两把相同的斧子），然后就是\(\frac{A^2(x)-B(x)}{2}\)，除2是因为顺序是无关紧要的
三个同理，\(\frac{A(x)^3-3A(x)B(x)+2C(x)}{6}\)
然后求和即可

代码

#include <cstdio>
#include <algorithm>
#include <cstring>
#define int long long
using namespace std;
const int MOD = 2281701377;
const int G = 3;
const int invG = 760567126;
const int MAXN = 300000;
int rev[MAXN];
int pow(int a,int b){
    int ans=1;
    while(b){
        if(b&1)
            ans=(1LL*ans*a)%MOD;
        a=(1LL*a*a)%MOD;
        b>>=1;
    }
    return ans;
}
void cal_rev(int n,int lim){
    for(int i=0;i<n;i++)
        rev[i]=(rev[i>>1]>>1)|((i&1)<<(lim-1));
}
void NTT(int *a,int opt,int n,int lim){
    for(int i=0;i<n;i++)
        if(i<rev[i])
            swap(a[i],a[rev[i]]);
    for(int i=2;i<=n;i<<=1){
        int len=i/2,tmp=pow((opt)?G:invG,(MOD-1)/i);
        for(int j=0;j<n;j+=i){
            int arr=1;
            for(int k=j;k<j+len;k++){
                int t=(1LL*a[k+len]*arr)%MOD;
                a[k+len]=(a[k]-t+MOD)%MOD;
                a[k]=(a[k]+t)%MOD;
                arr=(1LL*arr*tmp)%MOD;
            }
        }
    }
    if(!opt){
        int invN = pow(n,MOD-2);
        for(int i=0;i<n;i++)
            a[i]=(1LL*a[i]*invN)%MOD;
    }
}
int n,a[MAXN],b[MAXN],c[MAXN],d[MAXN],e[MAXN],f[MAXN],g[MAXN],h[MAXN],l[MAXN];
signed main(){
    scanf("%lld",&n);
    int maxx=0;
    for(int i=1;i<=n;i++){
        scanf("%lld",&a[i]);
        h[a[i]]++;
        b[a[i]]++;
        c[a[i]+a[i]]++;
        d[a[i]+a[i]+a[i]]++;        
        maxx=max(maxx,a[i]+a[i]+a[i]);
    }
    int midlen=1,midlim=0;
    while(midlen<=maxx)
        midlen<<=1,midlim++;
    cal_rev(midlen,midlim);
    NTT(b,1,midlen,midlim);
    for(int i=0;i<midlen;i++)
        l[i]=h[i];
    for(int i=0;i<midlen;i++)
        g[i]=(b[i]*b[i])%MOD;
    NTT(g,0,midlen,midlim);
    int INV6=pow(6,MOD-2);
    int INV2=pow(2,MOD-2);
    for(int i=0;i<midlen;i++)
        l[i]=(l[i]+(g[i]-c[i]+MOD)%MOD*INV2%MOD)%MOD;
    NTT(g,1,midlen,midlim);
    NTT(c,1,midlen,midlim);
    for(int i=0;i<midlen;i++)
        e[i]=(g[i]*b[i])%MOD,f[i]=(b[i]*c[i])%MOD;
    NTT(e,0,midlen,midlim);
    NTT(f,0,midlen,midlim);
    for(int i=0;i<midlen;i++)
        l[i]=l[i]+(e[i]-(3*f[i])%MOD+(2*d[i]))*INV6%MOD;
    // for(int i=0;i<midlen;i++)
    //     printf("e[%lld]=%lld\n",i,e[i]);
    // for(int i=0;i<midlen;i++)
    //     printf("f[%lld]=%lld\n",i,f[i]);
    for(int i=0;i<=maxx;i++)
        if(l[i])
            printf("%lld %lld\n",i,l[i]);
    return 0;
}