AGC 047C Product Modulo

题目描述

https://atcoder.jp/contests/agc047/tasks/agc047_c

简要题意：给定一个长度为 $n$ 的序列 $a_i$，求 $\sum_{i=1}^n\sum_{j=i+1}(a_i\times a_j\bmod P)$，注意只对乘积取模，$P=200003$

$n\le 2\times 10^5$

Solution

我们注意到如果不是只对成绩取模，那么显然可以直接 $FFT$，那么现在这种情况我们只能通过枚举值为 $x$ 和值为 $y$ 的乘积取模在乘上方案来计算答案

注意到 $P$ 是素数存在原根，那么我们转为枚举 $g^x$ 和 $g^y$，我们知道 $g^x\times g^y=g^{x+y}$ 这个东西符合卷积的形式，使用 $FFT$ 即可

时间复杂度 $O(n\log n)$

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#define ll long long
#define maxn 200010
using namespace std;

typedef std::vector<ll> Poly;
namespace Pol {
    //const int N = 4200000; // 200MB
    const int N = 2100000; // 100MB
    const double pi = acos(-1);
    struct Complex {
        double x, y;

        Complex(double x = 0, double y = 0) : x(x), y(y) {}

        friend Complex operator + (const Complex &u, const Complex &v) { return Complex(u.x + v.x, u.y + v.y); }
        friend Complex operator - (const Complex &u, const Complex &v) { return Complex(u.x - v.x, u.y - v.y); }
        friend Complex operator * (const Complex &u, const Complex &v) {
            return Complex(u.x * v.x - u.y * v.y, u.x * v.y + u.y * v.x); 
        }
    } a[N], b[N];
 
    int P[N];
    void init_P(int n) {
        int l = 0; while ((1 << l) < n) ++l;
        for (int i = 0; i < n; ++i) P[i] = (P[i >> 1] >> 1) | ((i & 1) << l - 1);
    }
    vector<Complex> init_W(int n) {
        vector<Complex> w(n); w[1] = 1;
        for (int i = 2; i < n; i <<= 1) {
            auto w0 = w.begin() + i / 2, w1 = w.begin() + i;
            Complex wn(cos(pi / i), sin(pi / i));
            for (int j = 0; j < i; j += 2)
                w1[j] = w0[j >> 1], w1[j + 1] = w1[j] * wn;
        }
        return w; 
    } auto w = init_W(1 << 21);
    void DIT(Complex *a, int n) {
        for (int k = n >> 1; k; k >>= 1)
            for (int i = 0; i < n; i += k << 1)
                for (int j = 0; j < k; ++j) {
                    Complex x = a[i + j], y = a[i + j + k];
                    a[i + j + k] = (x - y) * w[k + j], a[i + j] = x + y;
                }
    }
    void DIF(Complex *a, int n) {
        for (int k = 1; k < n; k <<= 1)
            for (int i = 0; i < n; i += k << 1)
                for (int j = 0; j < k; ++j) {
                    Complex x = a[i + j], y = a[i + j + k] * w[k + j];
                    a[i + j + k] = x - y, a[i + j] = x + y;
                }
        for (int i = 0; i < n; ++i) a[i].x /= n, a[i].y /= n;
        reverse(a + 1, a + n);
    }
    Poly Mul(const Poly &A, const Poly &B, int n1, int n2) {
        int n = 1; while (n < n1 + n2 - 1) n <<= 1; init_P(n);
        for (int i = 0; i < n1; ++i) a[i] = Complex(A[i], 0);
        for (int i = 0; i < n2; ++i) b[i] = Complex(B[i], 0);
        fill(a + n1, a + n, Complex(0, 0)); fill(b + n2, b + n, Complex(0, 0));
        DIT(a, n); DIT(b, n);
        for (int i = 0; i < n; ++i) a[i] = a[i] * b[i];
        DIF(a, n); Poly ans(n1 + n2 - 1); for (int i = 0; i < n1 + n2 - 1; ++i) ans[i] = (ll) (a[i].x + 0.5);
        return ans;
    }
}  // namespace Pol

const int p = 200003;
const int g = 2;
ll pow_mod(ll x, ll n) {
    ll s = 1;
    for (; n; n >>= 1, x = x * x % p)
        if (n & 1) s = s * x % p;
    return s;
}

int n, m, a[maxn];

int main() { 
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);

    cin >> n; Poly A(p); ll ans = 0; 
    for (int i = 1, x; i <= n; ++i) cin >> x, ++a[x];
    for (int i = 0, t = 1; i < p - 1; ++i) {
        A[i] = a[t];
        ans -= a[t] * (1ll * t * t % p);
        t = 1ll * t * g % p; 
    }
    Poly res = Pol::Mul(A, A, p, p);
    for (int i = 0; i < 2 * p - 1; ++i) 
        ans += 1ll * res[i] * pow_mod(g, i);
    cout << ans / 2 << "\n";
    return 0;
}