Primitive root + FFT + String Match

 

 

#include <stdio.h>
#include <math.h>
int root[100002]={...<Calculate 1e6 primitive root and put them here to get the root in O(1) ops>...};


long long A;
long long N,m,num[1000000];
long long ans[1000000],congruence[1000000],congruence_index[1000000], terms[1000000];
const double pi = acos (-1.0);

struct Complex {
    double r, i;
    Complex(double r = 0, double i = 0) : r(r), i(i) {
        //
    };
    Complex operator+ (const Complex &b) {
        return Complex (r + b.r, i + b.i);
    }
    Complex operator- (const Complex &b) {
        return Complex (r - b.r, i - b.i);
    }
    Complex operator* (const Complex &b) {
        return Complex (r * b.r - i * b.i, i * b.r + r * b.i);
    }
} s[1000000];

void reverse(Complex[], int);
void fft(Complex[], int);
void ifft(Complex[], int);
void _ifft(Complex[], int);
int main() {
//long long test = 1000000000000000000;
//    freopen("C:\\Users\\ASUS\\desktop\\1.in", "r", stdin);
//    freopen("C:\\Users\\ASUS\\desktop\\2.out","w",stdout);


    for (int i = 0; i < 1000000; i++) {
        num[i] = 0;
    }

    scanf("%lld%lld", &N, &m);

    if(m == 2){
        long long cnt = 0;
        for (int i = 1; i <= N; ++i) {
            scanf("%lld", &A);
            if(A % 2 == 1) cnt++;
        }
        printf("%lld\n", N * N - cnt * cnt);
        printf("%lld\n", cnt * cnt);
        return 0;
    }

    long long G = root[m];

    long long r = 1;
    for (int i = 1; i <= m - 1; ++i) {
        r = r * G % m;
        congruence[i] = r;
        congruence_index[r] = i;
    }

    //init


    for (int i = 0; i <= m; ++i) {
        ans[i] = 0;
    }
    //input


    for (int i = 1; i <= N; ++i) {
        scanf("%lld", &A);
        A %= m;
        if (A == 0) {
            continue;
        }
        num[congruence_index[A]] += 1;
    }

    //biggest term
    int n = 1;
    while (n <= 2 * m - 2) {
        n *= 2;
    }
    //Construct complex terms
    for (int i = 0; i < m; ++i) {
        s[i] = Complex(num[i], 0);
    }
    for (int i = m; i < n; ++i) {
        s[i] = Complex(0, 0);
    }
    fft(s, n);
    for (int i = 0; i < n; ++i) {
        s[i] = s[i] * s[i];
    }
    ifft(s, n);


//        for(int i = 0; i <= 2*m-2; ++i) {
//            long long now = (long long)(s[i].r + 0.5);
//            if(i%2==0) now -= num[i/2];
//            now/=2;
//            ans[mo[i%(m-1)]] += now;
//        }

    for (int i = 0; i <= 2 * m - 2; ++i) {
        terms[i] = (long long) (s[i].r + 0.5);
    }
    //addition
    for (int i = m; i <= 2 * m - 2; i++) {
        terms[i - (m - 1)] += terms[i];
    }
//    for(int i = 1; i < m; ++i){
//        printf("%d, %lld\n",i, terms[i]);
//    }
//    printf("\n\n");


    long long rest = 0;
    for (int i = 1; i < m; ++i) {
        ans[congruence[i]] = terms[i];
        rest += terms[i];

    }
    printf("%lld\n", (N * N - rest));

    for (int i = 1; i < m; ++i) {
        printf("%lld\n", ans[i]);
    }

//    for(int i = 0; i <= 2*m-2; ++i) {
//        long long now = (long long)(s[i].r + 0.5);
//        if(i%2==0) now -= num[i/2];
//        now/=2;
//        ans[congruence[i%(m-1)]] += now;
//    }
//        for(int i = 1; i < m; ++i) printf("%lld\n",ans[i]);

    return 0;
}


void ifft(Complex F[], int n){
    _ifft(F, n);
    for(int i = 0; i < n; i++){
        F[i].r /= n;
    }
}

void _ifft(Complex F[], int n){
    reverse(F, n);
    double t = -2.0 * pi;

    for(int h = 2; h <= n; h *= 2) {
        Complex wn (cos(t / h), sin(t / h));
        for(int j = 0; j < n; j += h) {
            Complex wk (1, 0);
            for(int k = j; k < j + h/2; ++k) {
                Complex u = F[k];
                Complex v = F[k + h/2] * wk;
                F[k]       = u + v;
                F[k + h/2] = u - v;
                wk = wk * wn;
            }
        }
    }
}

void fft(Complex F[], int n){
    reverse(F, n);
    double t = 2.0 * pi;

    for(int h = 2; h <= n; h *= 2) {
        Complex wn (cos(t / h), sin(t / h));
        for(int j = 0; j < n; j += h) {
            Complex wk (1, 0);
            for(int k = j; k < j + h/2; ++k) {
                Complex u = F[k];
                Complex v = F[k + h/2] * wk;
                F[k]       = u + v;
                F[k + h/2] = u - v;
                wk = wk * wn;
            }
        }
    }
//    for(int i = 0; i < n; i++){
//        printf("%f ", s[i].r);
//    }
//    printf("\n\n");
}

bool hasSwap[1000000];

void reverse(Complex a[], int len){
    int a_length = len;
    int n = (int)(log(len) / log(2));
    len = n;

    //init
    for(int i = 0; i < a_length; i++) hasSwap[i] = false;

    for(int i = 1; i < a_length-1; ++i) {
        if(hasSwap[i]) continue;
        int A = i;
        int j = 0;
        while(len-- > 0) {
            j <<= 1;
            if( (A & 1) == 1) {
                j += 1;
            }
            A >>= 1;
        }

        hasSwap[i] = true;
        hasSwap[j] = true;
        //swap
//        swap(a[i], a[j]);
        Complex tmp (a[i].r, a[i].i);
        a[i] = a[j];
        a[j] = tmp;
        len = n;
    }
}