Steins;Gate(原根+FFT)

http://newoj.acmclub.cn/problems/2072

题意：给出一个数组a，问对于每一个元素ak有多少个二元组(i,j)满足.

做法：求出p的原根g，然后上面的式子就变成如下形式

即求解有多少对bi+bj等于bk；

这个就好办首先我们队每一个数对P取模,然后得出他的指数，并对指数进行计数。

然后就用快速傅里叶变换，解决这个计数问题。

卷积后面的大于P的需要归入[0,P-1],直接取模循环就可以了。

注意0的情况特殊判断。

PS:在做题过程中原根求错了，抄了个红书的板子

#include "bits/stdc++.h"

using namespace std;
const double eps = 1e-6;
#define reg register
#define lowbit(x) x&-x
#define pll pair<ll,ll>
#define pii pair<int,int>
#define fi first
#define se second
#define makp make_pair
#define cp complex<double>

int dcmp(double x) {
    if (fabs(x) < eps) return 0;
    return (x > 0) ? 1 : -1;
}

typedef long long ll;
typedef unsigned long long ull;
const ull hash1 = 201326611;
const ull hash2 = 50331653;
const ll N = 1100000 + 10;
const int M = 1000000;
const int inf = 0x3f3f3f3f;
const ll mod = 1000000000 + 7;
const double Pi = acos(-1.0);
struct Complex {
    double x, y;

    Complex(double xx = 0, double yy = 0) { x = xx, y = yy; }
} a[N], omg[N], inv[N];
ll r[N], cnt = 0;

Complex operator+(Complex a, Complex b) { return Complex(a.x + b.x, a.y + b.y); }
Complex operator-(Complex a, Complex b) { return Complex(a.x - b.x, a.y - b.y); }
Complex operator*(Complex a, Complex b) { return Complex(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); }

void Cp_init(ll tot, ll lim) {
    for (int i = 0; i < tot; i++) {
        omg[i] = Complex(cos(2 * Pi * i / tot), sin(2 * Pi * i / tot));
        inv[i] = omg[i];
        inv[i].y *= -1;
    }
    for (int i = 0; i < tot; i++) {
        r[i] = (r[i >> 1] >> 1) | ((i & 1) << (lim - 1));
    }
}

void fft(Complex *a, Complex *omg, ll tot) {
    for (int i = 0; i < tot; i++) {
        if (i < r[i]) swap(a[i], a[r[i]]);
    }
    for (int l = 2; l <= tot; l <<= 1) {
        int m = l / 2;
        for (Complex *p = a; p != a + tot; p += l) {
            for (int i = 0; i < m; i++) {
                Complex tmp = omg[tot / l * i] * p[i + m];
                p[i + m] = p[i] - tmp;
                p[i] = p[i] + tmp;
            }
        }
    }
}

ll quick(ll a, ll n, ll p) {
    ll ans = 1;
    while (n) {
        if (n & 1) ans = ans * a % p;
        n >>= 1;
        a = a * a % p;
    }
    return ans;
}

ll vis[N], n, zero = 0;
ll num[N], xx[N], cont[N], I[N], p, ans[N];

vector<ll> gg;
bool g_test(ll g, ll p) {
    for (ll i = 0; i < gg.size(); i++) {
        if (quick(g, (p - 1) / gg[i], p) == 1)
            return 0;
    }
    return 1;
}

ll p_root(ll p) {
    ll tmp = p - 1;
    for (ll i = 2; i <= tmp / i; i++) {
        if (tmp % i == 0) {
            gg.push_back(i);
            while (tmp % i == 0) tmp /= i;
        }
    }
    if (tmp != 1) gg.push_back(tmp);
    ll g = 1;
    while (true) {
        if (g_test(g, p))
            return g;
        g++;
    }
}

int main() {
    scanf("%lld%lld", &n, &p);
    ll g = p_root(p);
    ///cout << g << endl;
    ll tt = 1;

    for (int i = 1; i < p; i++) {
        tt = tt * g % p;
        I[tt] = i;
    }
    I[1] = 0;

    for (int i = 0; i < n; i++) {
        scanf("%lld", &num[i]);
        xx[i] = num[i] % p;
        if (xx[i]) cont[I[xx[i]]]++;
        else zero++;
    }

    ll len1 = p - 1;
    ll lim = 0, tot = 1;

    while (tot < 2 * len1)
        tot <<= 1, lim++;

    for (int i = 0; i < len1; i++) a[i] = Complex(cont[i], 0);

    for (int i = len1; i < tot; i++) a[i] = Complex(0, 0);

    Cp_init(tot, lim);

    fft(a, omg, tot);

    for (int i = 0; i < tot; i++)
        a[i] = a[i] * a[i];

    fft(a, inv, tot);

    for (int i = 0; i < tot; i++)
        ans[i] = (ll) (a[i].x / tot + 0.5);

    for (int i = p - 1; i < tot; i++) ans[i % (p - 1)] += ans[i];

    ll Zans = 1LL * zero * (zero - 1) + 2LL * zero * (n - zero) + 1LL * zero;

    for (int i = 0; i < n; i++) {
        if (num[i] >= p) printf("0\n");
        else {
            if (xx[i]) printf("%lld\n", ans[I[xx[i]]]);
            else printf("%lld\n", Zans);
        }
    }
    return 0;
}