Problem: Let’s take a prime P = 200\,003. You are given N integers A_1, A_2, \ldots, A_N. Find the sum of ((A_i \cdot A_j) \bmod P) over all N \cdot (N-1) / 2 unordered pairs of elements (i < j).  Please note that the sum isn't computed modulo P.

Submission attempts (in order):
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  

@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
    
    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 = 210_000
FB = fft(B, fft_len)
C = np.rint(ifft(FB * FB, fft_len)).astype(np.int64)

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


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  

ntt_mod = 998_244_353
ntt_root = 3
max_ntt = 20
roots, iroots = np.empty(1 << max_ntt,
                         np.int64), np.empty(1 << max_ntt, np.int64)

ntt_data = (ntt_mod, ntt_root, max_ntt, roots, iroots)
ntt_t = numba.typeof(ntt_data)


@njit((ntt_t, ), cache=True)
def ntt_precompute(ntt_data):
    MOD, root, max_ntt, roots, iroots = ntt_data

    def mpow(a, n):
        p = 1
        while n:
            if n & 1:
                p = p * a % MOD
            a = a * a % MOD
            n >>= 1
        return p

    roots[0] = iroots[0] = 1
    for n in range(max_ntt):
        x = mpow(root, (MOD - 1) >> n + 2)
        y = mpow(x, MOD - 2)
        for i in range(1 << n):
            roots[(1 << n) | i] = roots[i] * x % MOD
            iroots[(1 << n) | i] = iroots[i] * y % MOD
    return roots, iroots


@njit((ntt_t, i8[:], b1), cache=True)
def ntt(ntt_data, A, inverse):
    MOD, _, _, roots, iroots = ntt_data

    def mpow(a, n):
        p = 1
        while n:
            if n & 1:
                p = p * a % MOD
            a = a * a % MOD
            n >>= 1
        return p

    N = len(A)

    if not inverse:
        m = N >> 1
        while m:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m] * roots[k]
                    A[i], A[i + m] = x + y, x - y
            A %= MOD
            m >>= 1
    else:
        m = 1
        while m < N:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m]
                    A[i], A[i + m] = x + y, (x - y) * iroots[k]
            A %= MOD
            m <<= 1
        for i in range(N):
            A[i] = A[i] * mpow(N, MOD - 2, MOD) % MOD
    return A


@njit((ntt_t, i8[:], i8[:]), cache=True)
def ntt_convolve(ntt_data, A, B):
    L = 1
    while 2 * L < len(A) + len(B) - 1:
        L *= 2
    L *= 2
    A = np.append(A, np.zeros(L - len(A), np.int64))
    B = np.append(B, np.zeros(L - len(B), np.int64))
    A = ntt(ntt_data, A, False)
    B = ntt(ntt_data, B, False)
    for i in range(len(A)):
        A[i] = A[i] * B[i] % ntt_data[0]
    return ntt(ntt_data, A, True)


ntt_precompute(ntt_data)

@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
    
    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)

B = np.append(B, [0] * ((1 << 19) - len(B)))
Bh, Bl = B >> 9, B & 511
x = ntt_convolve(ntt_data, Bh, Bh) << 18
x += ntt_convolve(ntt_data, Bh, Bl) << 10
x += ntt_convolve(ntt_data, Bl, Bl)

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

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  

@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
    
    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 = 420_000
FB = fft(B, fft_len)
C = np.rint(ifft(FB * FB, fft_len)).astype(np.int64)

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


Verdicts: Wrong Answer

Time Limit Exceeded

Accepted