Submission #64252544

---

Source Code Expand

#include<bits/stdc++.h>

using namespace std;

typedef long long i64;
constexpr double pi = acos(-1.0);

namespace Poly {

void FFT(vector<complex<double>> &s, int n, int o) {
    //
    vector<int> rev(n);
    for(int i = 1; i < n; i++) {
        rev[i] = (rev[i >> 1] >> 1) | ((i & 1) ? (n >> 1) : 0);
    }
    for(int i = 0; i < n; i++) {
        if(i < rev[i]) swap(s[i], s[rev[i]]);
    }
    //
    for(int i = 2; i <= n; i <<= 1) {
        complex<double> wn(cos(2 * pi / i), o * sin(2 * pi / i));
        for(int j = 0; j < n; j += i) {
            complex<double> wi(1, 0);
            for(int k = 0; k < (i >> 1); k++) {
                auto x = s[j + k], y = wi * s[j + (i >> 1) + k];
                s[j + k] = x + y, s[j + (i >> 1) + k] = x - y;
                wi *= wn;
            }
        }
    }
}

void DFT(vector<complex<double>> &s, int n) { FFT(s, n, 1); }
void IDFT(vector<complex<double>> &s, int n) { FFT(s, n, -1); }

}

int main() {
    cin.tie(nullptr) -> sync_with_stdio(false);

    int n;
    cin >> n;
    vector<int> s(n + 1);
    for(int i = 1; i <= n; i++) {
        cin >> s[i];
    }

    int w = *max_element(s.begin() + 1, s.end()); // w 次多项式
    int k = __lg(2 * w) + 1, len = (1 << k);
    vector<complex<double>> a(len);
    vector<int> cnt0(w + 1);
    for(int i = 1; i <= n; i++) {
        cnt0[s[i]]++;
    }
    for(int i = 1; i <= w; i++) {
        a[i].real(cnt0[i]);
    }

    Poly::DFT(a, len);
    for(int i = 0; i < len; i++) {
        a[i] *= a[i];
    }

    Poly::IDFT(a, len);
    vector<int> cnt(2 * w + 1);
    for(int i = 0; i <= 2 * w; i++) {
        int x = (int)round(a[i].real() / len);
        cnt[i] = x;
    }

    i64 ans = 0;
    for(int i = 1; i <= w; i++) {
        // x + y = 2 * i
        if(cnt0[i]) {
            i64 add = (cnt[2 * i] - 1) / 2;
            ans += add;
        }
    }
    cout << ans << '\n';
    
    return 0;
}

Submission Info

Judge Result