FFT+NTT 深入学习记录 2018秦皇岛ccpc-camp Steins;Gate 命运石之门

code

#define all(a) a.begin(),a.end()
#define pii pair<int,int>
#define pll pair<long long,long long>
#define mp(aa,bb) make_pair(aa,bb)
#define lrt rt<<1
#define rrt rt<<1|1
#define X first
#define Y second
#define PI (acos(-1.0))
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
const ll mod = 1000000007;
const double eps=1e-9;
const int inf=0x3f3f3f3f;
typedef long long LL;
const int maxn = 2e6+10;
// head  
// meopass 's mo ban
LL a[maxn], to1[maxn], to2[maxn];
LL f1[maxn], f2[maxn];
LL N,P;
LL Pow(LL a, LL b, LL m) {
    LL res = 1;
    while(b) {
        if(b & 1) res = (res * a) % m;
        a = (a * a) % m;
        b >>= 1;
    }
    return res;
}
  
LL Root(LL m) {
    vector<LL> V;
    LL buf = m - 1;
    for(LL i = 2; i <= sqrt(buf); i++) {
        if(buf % i == 0) {
            V.push_back(i);
            while(buf % i == 0) buf /= i;
        }
    }
    if(buf != 1) V.push_back(buf);
    for(LL i = 2; ; i++) {
        LL j = 0;
        for(; j < V.size() && Pow(i, (m-1)/V[j], m) != 1; j++) ;
        if(j == V.size()) return i;
    }
}
  
typedef long long LL;
const LL kMaxn = 8e5+20;
typedef long double ld;
const double kPi = acos(-1.0);
LL len, n;
  
  
struct Complex {
    double r, i;
    Complex(double r = 0, double i = 0):r(r), i(i) {};
    Complex operator + (const Complex & rhs) {
        return Complex(r + rhs.r, i + rhs.i);
    }
    Complex operator - (const Complex & rhs) {
        return Complex(r - rhs.r, i - rhs.i);
    }
    Complex operator * (const Complex &rhs) {
        return Complex(r * rhs.r - i * rhs.i, i * rhs.r + r * rhs.i);
    }
} va[kMaxn], vb[kMaxn];
  
void Rader(Complex F[], LL len) {
    LL j = len >> 1;
    for(LL i = 1; i < len - 1; i++) {
        if(i < j) swap(F[i], F[j]);
        LL k = len >> 1;
        while(j >= k) {
            j -= k;
            k >>= 1;
        }
        if(j < k) j += k;
    }
}
  
void FFT(Complex F[], LL len, LL t) { //时域转频域
    Rader(F, len);
    for(LL h = 2; h <= len; h <<= 1) {
        Complex wn(cos(-t*2*kPi/h), sin(-t*2*kPi/h));
        for(LL j = 0; j < len; j += h) {
            Complex E(1, 0);
            for(LL k = j; k < j + h / 2; k++) {
                Complex u = F[k];
                Complex v = E*F[k+h/2];
                F[k] = u + v;
                F[k+h/2] = u - v;
                E = E * wn;
            }
        }
    }
    if(t == -1) {
        for(LL i = 0; i < len; i++) {
            F[i].r /= len;
        }
    }
}
  
void Init(LL *a, LL *b) {
    LL n1 = 2e5+10, n2 = 2e5+10;
    len = 1;
    while(len < 2*n1 || len < 2*n2) len <<= 1;
    LL i;
    for(i = 0; i < n1; i++) {
        va[i].r = a[i], va[i].i = 0;
    }
    while(i < len) va[i].r = va[i].i = 0, i++;
    for(i = 0; i < n2; i++) {
        vb[i].r = b[i], vb[i].i = 0;
    }
    while(i < len) vb[i].r = vb[i].i = 0, i++;
}
  
void Solve(LL *a, LL *b) {
    Init(a, b);
    FFT(va, len, 1);
    FFT(vb, len, 1);
    for(LL i = 0; i < len; i++) va[i] = va[i] * vb[i];
    FFT(va, len, -1);
    for(LL i = 0; i < len; i++)
        a[i] = va[i].r + 0.5;
}
  
 
// meopass 's ban zi end
map<LL,vector<LL>>mp;
LL ans[maxn];
int main()
{   
    sldd(N,P);
    for(int i=1;i<=N;i++)
    {
        LL t;
        sld(t);
        mp[t].push_back(i);
        a[t%P]++;
    }
    //  find the root 
    LL R = Root(P);
    LL tmp=1;
    for(LL i=0;i<P-1;i++)
    {
        to1[tmp] = i;
        to2[i]=tmp;
        tmp=R*tmp%P;
    }
    for(LL i=0;i<P-1;i++)
    {
        f1[i]=a[to2[i]];
        f2[i]=f1[i];
    }
    Solve(f1,f2);
    for(LL i=0;i<P-1;i++)
        f2[i]=a[to2[i]];
    for(auto meo:mp[0])
        ans[meo]=a[0]*N*2-a[0]*a[0];
    for(LL i=1;i<P;i++)
    {
        for(auto meo:mp[i])
          ans[meo]=f1[to1[i]]+f1[to1[i]+P-1];
    }
    for(LL i=1;i<=N;i++)
        printf("%lld\n",ans[i]);
    return 0;
}