結果

問題 No.931 Multiplicative Convolution
ユーザー tktk_snsn
提出日時 2021-06-03 20:36:22
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 732 ms / 2,000 ms
コード長 1,551 bytes
コンパイル時間 123 ms
コンパイル使用メモリ 12,672 KB
実行使用メモリ 69,196 KB
最終ジャッジ日時 2024-11-17 08:24:54
合計ジャッジ時間 15,019 ms
ジャッジサーバーID
(参考情報)
judge2 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 14
権限があれば一括ダウンロードができます

ソースコード

diff #

import numpy as np
fft = np.fft.rfft
ifft = np.fft.irfft
mod = 998244353


def gen(p):
    fact = []
    phi = p - 1
    n = phi
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0:
            fact.append(i)
            while n % i == 0:
                n //= i
    if n > 1:
        fact.append(n)
    for res in range(2, p+1):
        ok = True
        for f in fact:
            if pow(res, phi//f, p) == 1:
                ok = False
                break
        if ok:
            return res
    return -1


def conv(A, B):
    n = len(A) + len(B) - 1
    sz = 1 << (n-1).bit_length()
    res = ifft(fft(A, sz) * fft(B, sz), sz)
    return np.rint(res).astype(np.int64)[:n]


def mod_conv(A, B):
    a1, a2 = np.divmod(A, 1 << 15)
    b1, b2 = np.divmod(B, 1 << 15)
    x = conv(a1, b1) % mod
    y = conv(a2, b2) % mod
    xy = (conv(a1 + a2, b1 + b2) - (x + y)) % mod
    res = (x << 30) % mod + (xy << 15) % mod + y
    return res % mod


def mod_pow(n, power, mod):
    """ [n^0, n^1, n^2, ..., n^(power-1)] """
    D = power.bit_length()
    res = np.empty(1 << D, np.int64)
    res[0] = 1
    for d in range(D):
        res[(1 << d): (1 << (d + 1))] = \
            res[:(1 << d)] * n % mod * res[(1 << d) - 1] % mod
    return res[:power]


P = int(input())
A = np.array(input().split(), dtype=np.int64)
B = np.array(input().split(), dtype=np.int64)
g = gen(P)
idx = mod_pow(g, P - 1, P) - 1

X = A[idx]
Y = B[idx]
Z = mod_conv(X, Y)

Z[:P-2] += Z[P-1:]

ans = np.empty(P-1, dtype=np.int64)
ans[idx] = Z[:P-1] % mod
print(*ans)
0