提出 #15780907

---

ソースコード 拡げる

import sys
import numpy as np
import numba
from numba import njit, b1, i4, i8, f8

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

MOD = 200_003
R = 2  # primitive root

@njit((i8[:], ), cache=True)
def precompute(A):
    exp = np.zeros(MOD - 1, np.int64)
    log = np.zeros(MOD, np.int64)
    exp[0] = 1
    for k in range(1, MOD - 1):
        exp[k] = exp[k - 1] * R % MOD
        log[exp[k]] = k
    # r^k の個数を B[k] に入れる
    B = np.zeros(MOD - 1, np.int64)
    for x in A:
        if x != 0:
            B[log[x]] += 1
    return exp, log, B

@njit((i8[:], i8[:], i8[:], i8[:]), cache=True)
def get_ans(A, C, exp, log):
    # 順序対について集計
    ans = 0
    for i in range(len(C)):
        x = exp[i % (MOD - 1)]
        ans += x * C[i]
    for a in A:
        x = a * a % MOD
        ans -= x
    return ans // 2

A = np.array(read().split(), np.int64)[1:]

exp, log, B = precompute(A)
fft, ifft = np.fft.fft, np.fft.ifft
fft_len = 1 << 20
FB = fft(B, fft_len)
C = np.rint(ifft(FB * FB, fft_len)).astype(np.int64)

print(get_ans(A, C, exp, log))

提出情報

ジャッジ結果