AtCoder Grand Contest 047 C

C - Product Modulo

問題

$P = 2 \cdot 10^5 + 3$ (素数)とする。

$n$ 個の整数 $\{ a_i \}$ が与えられる。 $\sum_{i \lt j} (a_i a_j \bmod P)$ を求めよ。

制約

$2 \leq n \leq 2 \cdot 10^5$
$0 \leq a_i \lt P$

考察

コンテスト中は $x \bmod P = x - \lfloor x / P \rfloor \cdot P$ を使うとずっと思っていたが違った。

以降 $a_i = 0$ なる $i$ は無視して考える。

$g = 2$ は $P$ の原始根なので、任意の $1 \leq x \lt P$ に対して $g^k \equiv x$ なる $0 \leq k \leq P - 2$ が存在する。そこで、 $g^{b_i} = a_i$ を満たす数列 $\{ b_i \}$ を考える。すると、 $a_i a_j \equiv g^{b_i + b_j}$ となり、「 $\{a_i\}$ 同士の掛け算」が「 $\{b_i\}$ 同士の足し算」に変わる。よって各 $0 \leq k \leq 2(P - 2)$ について $b_i + b_j = k$ なる $i \lt j$ を数え上げればよく、これは FFT によって $\Theta(P \log P)$ でできる。

実装例

提出 #15794827 - AtCoder Grand Contest 047

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56


#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
#include <complex>

template <int MOD>
struct ModInt { ... };

template <int K>
struct FastFourierTransform { ... };

constexpr int MOD = 200003;
using mint = ModInt<MOD>;
const FastFourierTransform<20> FFT;

using lint = long long;

void solve() {
    // g^log[x] = x
    std::vector<int> log(MOD);
    mint g = 2;
    for (int i = 0; i < MOD - 1; ++i) {
        log[g.pow(i).val] = i;
    }

    std::vector<lint> cnt(MOD - 1, 0);

    int n;
    std::cin >> n;
    while (n--) {
        int a;
        std::cin >> a;
        if (a != 0) ++cnt[log[a]];
    }

    // FFTで畳み込む
    auto res = FFT.fft(cnt, cnt);

    lint ans = 0;
    for (int i = 0; i < (int)res.size(); ++i) {
        lint num = llround(res[i]);
        ans += g.pow(i).val * num;
    }

    // a_i * a_iを省く
    for (int i = 0; i < (int)cnt.size(); ++i) {
        lint num = llround(cnt[i]);
        ans -= g.pow(i * 2).val * num;
    }

    // i > jを省く
    ans /= 2;

    std::cout << ans << "\n";
}