結果

問題 No.2994 べき内積
ユーザー むつある
提出日時 2024-12-19 01:53:22
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 4,693 bytes
コンパイル時間 434 ms
コンパイル使用メモリ 82,392 KB
実行使用メモリ 169,800 KB
最終ジャッジ日時 2024-12-19 01:54:24
合計ジャッジ時間 62,278 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1 WA * 2
other AC * 1 WA * 2 TLE * 20
権限があれば一括ダウンロードができます

ソースコード

diff # 

mod = 1009

# ---- NTTのための前処理 ----
# 原始根g=5を利用。（1009でのprimitive rootは5が一例）
g = 5

# p-1 = 1008 = 2^4 * 63
# よって最大で長さ 16 (=2^4)を基本としたNTTが可能。
# しかし多項式畳み込みにはより長い長さが必要になる場合があるため、
# ここでは必要な長さに応じてNTTを実行する。
# 必要な長さを2べきに拡張し、その2べきに対する原始根の(p-1)/length乗を計算する。

def modexp(base, exp, m):
    res = 1
    cur = base % m
    while exp > 0:
        if exp & 1:
            res = (res * cur) % m
        cur = (cur * cur) % m
        exp >>= 1
    return res

def ntt(a, inv=False):
    """
    inplaceでNTTまたはINTTを行う
    a: 配列（長さは2冪）
    inv: FalseならNTT、Trueなら逆NTT
    """
    n = len(a)
    # ビット反転
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j |= bit
        if i < j:
            a[i], a[j] = a[j], a[i]
    
    # nの最大長に応じて原始根の適切な冪を求める
    # root = g^{(1008/n)} mod p
    # ここでp=1009, p-1=1008
    # NTT (forward): root = g^{(1008/n)}
    # INTT (inverse): root = (g^{(1008/n)})^{-1}
    length = 1
    if inv:
        root = modexp(g, 1008 // n, mod)
        root_inv = modexp(root, mod-2, mod)
        # 逆変換用はroot_invを使う
    else:
        root = modexp(g, 1008 // n, mod)
    
    step = 1
    while step < n:
        step2 = step << 1
        if inv:
            w = 1
            for i in range(0, n, step2):
                for j in range(i, i+step):
                    u = a[j]
                    v = a[j+step]*w % mod
                    a[j] = (u+v)%mod
                    a[j+step] = (u-v)%mod
                w = (w*root_inv) % mod
        else:
            w = 1
            for i in range(0, n, step2):
                for j in range(i, i+step):
                    u = a[j]
                    v = a[j+step]*w % mod
                    a[j] = (u+v)%mod
                    a[j+step] = (u-v)%mod
                w = (w*root) % mod
        step = step2
    
    if inv:
        # nで割る
        ninv = modexp(n, mod-2, mod)
        for i in range(n):
            a[i] = a[i]*ninv % mod

def polymul(a, b):
    """
    多項式a,bの積をNTTで計算
    a,bはmod 1009下の係数を持つ多項式とする
    """
    n = len(a)
    m = len(b)
    size = 1
    while size < n+m-1:
        size <<= 1
    A = a[:] + [0]*(size-n)
    B = b[:] + [0]*(size-m)
    
    ntt(A, inv=False)
    ntt(B, inv=False)
    for i in range(size):
        A[i] = (A[i]*B[i])%mod
    ntt(A, inv=True)
    
    # 結果は長さ n+m-1 の多項式
    return A[:n+m-1]

def toeplitz_multiply(a, b):
    """
    a,bは長さNの下三角トープリッツ行列A,Bの1列目。
    c = A*B の1列目を返す。
    c_n = sum_{m=0}^n a_m b_{n-m}  (n=0 to N-1)
    多項式 A(x)=sum a_m x^m, B(x)=sum b_m x^m
    A(x)*B(x) の先頭N項が c となる。
    """
    # 多項式積
    c_full = polymul(a, b)
    # 先頭N項が必要
    return c_full[:len(a)]

def toeplitz_identity(N):
    """
    N×Nの単位行列を下三角トープリッツ形式で表すと、
    第一列は [1, 0, 0, ..., 0]
    """
    a = [0]*N
    a[0] = 1
    return a

def toeplitz_power(a, K):
    """
    下三角トープリッツ行列 A (第一列 a) の K 乗 A^K を求める。
    二分累乗法を用いる。
    """
    N = len(a)
    # 結果を単位行列で初期化
    result = toeplitz_identity(N)
    base = a[:]

    k = K
    while k > 0:
        if k & 1:
            result = toeplitz_multiply(result, base)
        base = toeplitz_multiply(base, base)
        k >>= 1
    
    return result
    

M,N=map(int,input().split())
M+=1
N+=1
k=list(map(int,input().split()))
a=list(map(int,input().split()))

b=[[a[i]] for i in range(N)]
p=1009

for i in range(M):
  if k[i]!=0:
    x=i
    k[x]-=1
    break
  for j in range(x):
    k[j]=p-1

keep=[]

for i in range(M):
  keep.append(toeplitz_power(a, k[i]))
  a=toeplitz_power(a, p)

now=toeplitz_identity(N)
for i in range(M):
  now=toeplitz_multiply(now, keep[i])

x=[[0]*N for _ in range(N)]
for i in range(N):
  for j in range(i+1):
    x[i][j]=now[i-j]

def mat_mul(a, b) :
    I, J, K = len(a), len(b[0]), len(b)
    c = [[0] * J for _ in range(I)]
    for i in range(I) :
        for j in range(J) :
            for k in range(K) :
                c[i][j] += a[i][k] * b[k][j]
                c[i][j] %= p
    return c
    
ans=mat_mul(x,b)
for i in ans:
  print(i[0]%p,end=' ')
print()
0