結果

問題 No.2506 Sum of Weighted Powers
ユーザー suisensuisen
提出日時 2023-03-22 01:54:37
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,185 ms / 2,000 ms
コード長 6,243 bytes
コンパイル時間 364 ms
コンパイル使用メモリ 81,920 KB
実行使用メモリ 164,476 KB
最終ジャッジ日時 2024-09-15 14:11:52
合計ジャッジ時間 14,399 ms
ジャッジサーバーID
(参考情報)
judge1 / judge3
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 71 ms
68,096 KB
testcase_01 AC 73 ms
67,840 KB
testcase_02 AC 74 ms
68,096 KB
testcase_03 AC 117 ms
79,104 KB
testcase_04 AC 81 ms
70,400 KB
testcase_05 AC 75 ms
68,992 KB
testcase_06 AC 125 ms
79,360 KB
testcase_07 AC 118 ms
79,104 KB
testcase_08 AC 127 ms
79,488 KB
testcase_09 AC 120 ms
78,592 KB
testcase_10 AC 126 ms
79,616 KB
testcase_11 AC 126 ms
78,848 KB
testcase_12 AC 76 ms
68,736 KB
testcase_13 AC 105 ms
78,464 KB
testcase_14 AC 79 ms
70,144 KB
testcase_15 AC 79 ms
70,016 KB
testcase_16 AC 105 ms
78,720 KB
testcase_17 AC 120 ms
79,104 KB
testcase_18 AC 120 ms
79,232 KB
testcase_19 AC 106 ms
78,080 KB
testcase_20 AC 123 ms
78,720 KB
testcase_21 AC 125 ms
79,488 KB
testcase_22 AC 119 ms
78,592 KB
testcase_23 AC 730 ms
126,176 KB
testcase_24 AC 927 ms
134,420 KB
testcase_25 AC 388 ms
102,260 KB
testcase_26 AC 1,010 ms
140,652 KB
testcase_27 AC 196 ms
82,004 KB
testcase_28 AC 1,169 ms
164,476 KB
testcase_29 AC 1,181 ms
163,844 KB
testcase_30 AC 1,178 ms
163,828 KB
testcase_31 AC 1,168 ms
164,476 KB
testcase_32 AC 1,185 ms
163,720 KB
testcase_33 AC 74 ms
68,096 KB
testcase_34 AC 71 ms
68,096 KB
testcase_35 AC 165 ms
124,560 KB
testcase_36 AC 181 ms
132,808 KB
testcase_37 AC 71 ms
68,096 KB
testcase_38 AC 72 ms
67,968 KB
testcase_39 AC 73 ms
68,096 KB
testcase_40 AC 72 ms
67,968 KB
testcase_41 AC 71 ms
68,096 KB
権限があれば一括ダウンロードができます

ソースコード

diff #

from typing import List

def bsf(x):
    res = 0
    while not (x & 1):
        res += 1
        x >>= 1
    return res

P = 943718401
G = 7
rank2 = bsf(P - 1)

class NTT:
    class __RootInitializer:
        @staticmethod
        def root():
            return [pow(G, (P - 1) >> i, P) for i in range(0, rank2 + 1)]

        @staticmethod
        def iroot():
            return [pow(pow(G, P - 2, P), (P - 1) >> i, P) for i in range(0, rank2 + 1)]

    root = __RootInitializer.root()
    iroot = __RootInitializer.iroot()

    class __RateInitializer:
        @staticmethod
        def rates(root: List[int], iroot: List[int]):
            rate2 = [0] * max(0, rank2 - 1)
            irate2 = [0] * max(0, rank2 - 1)
            prod = iprod = 1
            for i in range(rank2 - 1):
                rate2[i] = root[i + 2] * prod % P
                irate2[i] = iroot[i + 2] * iprod % P
                prod = prod * iroot[i + 2] % P
                iprod = iprod * root[i + 2] % P
            
            rate3 = [0] * max(0, rank2 - 2)
            irate3 = [0] * max(0, rank2 - 2)
            prod = iprod = 1
            for i in range(rank2 - 2):
                rate3[i] = root[i + 3] * prod % P
                irate3[i] = iroot[i + 3] * iprod % P
                prod = prod * iroot[i + 3] % P
                iprod = iprod * root[i + 3] % P
            return rate2, irate2, rate3, irate3
    
    rate2, irate2, rate3, irate3 = __RateInitializer.rates(__RootInitializer.root(), __RootInitializer.iroot())

    @staticmethod
    def butterfly(a: List[int]) -> None:
        n = len(a)
        h = bsf(n)
        l = 0
        while l < h:
            if h - l == 1:
                p = 1 << (h - l - 1)
                rot = 1
                for s in range(1 << l):
                    offset = s << (h - l)
                    for i in range(p):
                        u = a[i + offset]
                        v = a[i + offset + p] * rot
                        a[i + offset] = (u + v) % P
                        a[i + offset + p] = (u - v) % P
                    if s + 1 != 1 << l:
                        rot = rot * NTT.rate2[bsf(~s)] % P
                l += 1
            else:
                p = 1 << (h - l - 2)
                rot, imag = 1, NTT.root[2]
                for s in range(1 << l):
                    rot2 = rot * rot % P
                    rot3 = rot2 * rot % P
                    offset = s << (h - l)
                    for i in range(p):
                        a0 = a[i + offset]
                        a1 = a[i + offset + p] * rot
                        a2 = a[i + offset + 2 * p] * rot2
                        a3 = a[i + offset + 3 * p] * rot3
                        a1na3imag = (a1 - a3) % P * imag
                        a[i + offset] = (a0 + a2 + a1 + a3) % P
                        a[i + offset + 1 * p] = (a0 + a2 - a1 - a3) % P
                        a[i + offset + 2 * p] = (a0 - a2 + a1na3imag) % P
                        a[i + offset + 3 * p] = (a0 - a2 - a1na3imag) % P
                    if s + 1 != (1 << l):
                        rot = rot * NTT.rate3[bsf(~s)] % P
                l += 2

    @staticmethod
    def butterfly_inv(a : List[int]) -> None:
        n = len(a)
        h = bsf(n)

        l = h
        while l:
            if l == 1:
                p = 1 << (h - l)
                irot = 1
                for s in range(1 << (l - 1)):
                    offset = s << (h - l + 1)
                    for i in range(p):
                        u = a[i + offset]
                        v = a[i + offset + p]
                        a[i + offset] = (u + v) % P
                        a[i + offset + p] = ((u - v) * irot) % P
                    if s + 1 != 1 << (l - 1):
                        irot = irot * NTT.irate2[bsf(~s)] % P
                l -= 1
            else:
                p = 1 << (h - l)
                irot = 1
                iimag = NTT.iroot[2]
                for s in range(1 << (l - 2)):
                    irot2 = irot * irot % P
                    irot3 = irot2 * irot % P
                    offset = s << (h - l + 2)
                    for i in range(p):
                        a0 = a[i + offset]
                        a1 = a[i + offset + p]
                        a2 = a[i + offset + 2 * p]
                        a3 = a[i + offset + 3 * p]

                        a2na3iimag = (a2 - a3) * iimag % P

                        a[i + offset] = (a0 + a1 + a2 + a3) % P
                        a[i + offset + p] = ((a0 - a1 + a2na3iimag) * irot) % P
                        a[i + offset + 2 * p] = ((a0 + a1 - a2 - a3) * irot2) % P
                        a[i + offset + 3 * p] = ((a0 - a1 - a2na3iimag) * irot3) % P
                    if s + 1 != 1 << (l - 2):
                        irot = irot * NTT.irate3[bsf(~s)] % P
                l -= 2
    
    @staticmethod
    def convolution(a, b):
        n = len(a)
        m = len(b)
        if not a or not b:
            return []
        if min(n, m) <= 40:
            if n < m:
                n, m = m, n
                a, b = b, a
            res = [0] * (n + m - 1)
            for i in range(n):
                for j in range(m):
                    res[i + j] += a[i] * b[j]
                    res[i + j] %= P
            return res
        z = 1 << ((n + m - 1).bit_length())

        iz = pow(z, P - 2, P)
        a += [0] * (z - n)
        b += [0] * (z - m)
        NTT.butterfly(a)
        NTT.butterfly(b)
        c = [a[i] * b[i] % P * iz % P for i in range(z)]
        NTT.butterfly_inv(c)
        return c[:n + m - 1]

n, x = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
C = list(map(int, input().split()))

if x == 0:
    ans = sum(A[i] * B[i] % P * C[0] % P for i in range(n + 1)) + sum(A[i] * B[0] % P * C[i] % P for i in range(1, n + 1))
    print(ans % P)
else:
    t = lambda k : (k - 1) * k * (k + 1) // 3

    inv_x = pow(x, P - 2, P)

    F = [B[i] * pow(inv_x, t(i), P) for i in range(n + 1)]
    G = [C[i] * pow(inv_x, t(i), P) for i in range(n + 1)]

    H = NTT.convolution(F, G)

    ans = sum(A[i] * pow(x, t(i), P) % P * H[i] % P for i in range(n + 1))

    print(ans % P)
0