結果

問題 No.931 Multiplicative Convolution
ユーザー maspymaspy
提出日時 2023-05-12 21:12:42
言語 Python3
(3.12.2 + numpy 1.26.4 + scipy 1.12.0)
結果
AC  
実行時間 812 ms / 2,000 ms
コード長 2,145 bytes
コンパイル時間 432 ms
コンパイル使用メモリ 12,928 KB
実行使用メモリ 71,148 KB
最終ジャッジ日時 2024-11-28 16:53:07
合計ジャッジ時間 14,525 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 506 ms
44,112 KB
testcase_01 AC 511 ms
44,364 KB
testcase_02 AC 565 ms
44,616 KB
testcase_03 AC 511 ms
44,240 KB
testcase_04 AC 515 ms
44,752 KB
testcase_05 AC 509 ms
44,108 KB
testcase_06 AC 514 ms
44,292 KB
testcase_07 AC 535 ms
45,008 KB
testcase_08 AC 775 ms
71,148 KB
testcase_09 AC 704 ms
65,888 KB
testcase_10 AC 802 ms
68,368 KB
testcase_11 AC 705 ms
67,356 KB
testcase_12 AC 717 ms
58,832 KB
testcase_13 AC 812 ms
66,032 KB
testcase_14 AC 778 ms
70,868 KB
testcase_15 AC 771 ms
70,840 KB
testcase_16 AC 772 ms
67,068 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

# 原子根のべき乗で書いて、多項式に帰着

import numpy as np

p = int(readline())
A = np.array([0] + readline().split(), np.int64)
B = np.array([0] + readline().split(), np.int64)

MOD = 998244353

def make_prime(U):
    is_prime = np.zeros(U,np.bool_)
    is_prime[2] = 1
    is_prime[3::2] = 1
    M = int(U**.5)+1
    for p in range(3,M,2):
        if is_prime[p]:
            is_prime[p*p::p+p] = 0
    return is_prime, is_prime.nonzero()[0]

def find_primitive_root(p):
    import random
    if p == 2:
        return 1
    _, primes = make_prime(p)
    e = p - 1
    div = [q for q in primes if not e % q]
    test_n = [e // q for q in div]
    while True:
        r = random.randint(1,p-1)
        if any(pow(r,int(n),p) == 1 for n in test_n):
            continue
        return r

def make_power(a, L, MOD=MOD):
    B = L.bit_length()
    x = np.empty((1<<B), np.int64)
    x[0] = 1
    for n in range(B):
        x[1<<n:1<<(n+1)] = x[:1<<n] * a % MOD
        a *= a; a %= MOD
    return x[:L]

r = find_primitive_root(p)
exp = make_power(r,p-1,p)
log = np.zeros(p,np.int32)
log[exp] = np.arange(p-1)

f = A[exp]; g = B[exp]

def fft_convolve(f, g, MOD = MOD):
    """
    数列 (多項式) f, g の畳み込みの計算．上下 15 bitずつ分けて計算することで，
    30 bit以下の整数，長さ 250000 程度の数列での計算が正確に行える．
    """
    fft = np.fft.rfft; ifft = np.fft.irfft
    Lf = len(f); Lg = len(g); L = Lf + Lg - 1
    fft_len = 1 << L.bit_length()
    fl = f & (1 << 15) - 1; fh = f >> 15
    gl = g & (1 << 15) - 1; gh = g >> 15
    conv = lambda f,g: ifft(fft(f,fft_len) * fft(g,fft_len))[:L]
    x = conv(fl, gl) % MOD
    y = conv(fl+fh, gl+gh) % MOD
    z = conv(fh, gh) % MOD
    a, b, c = map(lambda x: (x + .5).astype(np.int64), [x,y,z])
    return (a + ((b - a - c) << 15) + (c << 30)) % MOD

fg = fft_convolve(f,g)

L = len(fg); fg[0:L-(p-1)] += fg[p-1:L]; fg = fg[:p-1]; fg %= MOD

answer = fg[log[1:]]
print(' '.join(answer.astype(str)))
0