#944949 (PyPy3) No.2613 Sum of Combination

提出ソース

結果

問題	No.2613 Sum of Combination
ユーザー	chineristAC
提出日時	2024-01-19 21:51:26
言語	PyPy3 (7.3.15)
結果	AC
実行時間	877 ms / 4,500 ms
コード長	8,188 bytes
コンパイル時間	174 ms
コンパイル使用メモリ	81,700 KB
実行使用メモリ	159,420 KB
最終ジャッジ日時	2024-01-19 21:51:53
合計ジャッジ時間	24,814 ms
ジャッジサーバーID （参考情報）	judge13 / judge15

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	64 ms 68,928 KB
testcase_01	AC	64 ms 68,928 KB
testcase_02	AC	151 ms 84,076 KB
testcase_03	AC	64 ms 68,928 KB
testcase_04	AC	65 ms 68,928 KB
testcase_05	AC	65 ms 68,928 KB
testcase_06	AC	64 ms 68,928 KB
testcase_07	AC	66 ms 68,928 KB
testcase_08	AC	123 ms 82,868 KB
testcase_09	AC	73 ms 73,284 KB
testcase_10	AC	64 ms 68,928 KB
testcase_11	AC	105 ms 82,100 KB
testcase_12	AC	102 ms 82,100 KB
testcase_13	AC	156 ms 84,216 KB
testcase_14	AC	156 ms 84,240 KB
testcase_15	AC	156 ms 83,464 KB
testcase_16	AC	154 ms 84,200 KB
testcase_17	AC	157 ms 84,276 KB
testcase_18	AC	151 ms 83,796 KB
testcase_19	AC	158 ms 84,488 KB
testcase_20	AC	116 ms 82,236 KB
testcase_21	AC	94 ms 79,644 KB
testcase_22	AC	185 ms 87,884 KB
testcase_23	AC	810 ms 150,252 KB
testcase_24	AC	780 ms 150,008 KB
testcase_25	AC	784 ms 143,900 KB
testcase_26	AC	818 ms 158,472 KB
testcase_27	AC	488 ms 121,912 KB
testcase_28	AC	838 ms 158,064 KB
testcase_29	AC	800 ms 154,808 KB
testcase_30	AC	851 ms 158,296 KB
testcase_31	AC	804 ms 152,140 KB
testcase_32	AC	794 ms 151,404 KB
testcase_33	AC	850 ms 158,520 KB
testcase_34	AC	799 ms 157,752 KB
testcase_35	AC	822 ms 159,420 KB
testcase_36	AC	820 ms 159,420 KB
testcase_37	AC	816 ms 158,780 KB
testcase_38	AC	832 ms 157,224 KB
testcase_39	AC	832 ms 157,608 KB
testcase_40	AC	805 ms 158,124 KB
testcase_41	AC	835 ms 158,520 KB
testcase_42	AC	830 ms 157,992 KB
testcase_43	AC	850 ms 159,276 KB
testcase_44	AC	817 ms 157,092 KB
testcase_45	AC	66 ms 68,928 KB
testcase_46	AC	66 ms 68,928 KB
testcase_47	AC	66 ms 68,928 KB
testcase_48	AC	67 ms 71,056 KB
testcase_49	AC	71 ms 71,056 KB
testcase_50	AC	877 ms 158,240 KB
testcase_51	AC	810 ms 157,092 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

import sys
from itertools import permutations
from heapq import heappop,heappush
from collections import deque
import random
import bisect
from math import gcd


input = lambda :sys.stdin.readline().rstrip()
mi = lambda :map(int,input().split())
li = lambda :list(mi())

mod = 998244353
omega = pow(3,119,mod)
rev_omega = pow(omega,mod-2,mod)

N = 2*10**5
g1 = [1]*(N+1) # 元テーブル
g2 = [1]*(N+1) #逆元テーブル
inv = [1]*(N+1) #逆元テーブル計算用テーブル

for i in range( 2, N + 1 ):
    g1[i]=( ( g1[i-1] * i ) % mod )
    inv[i]=( ( -inv[mod % i] * (mod//i) ) % mod )
    g2[i]=( (g2[i-1] * inv[i]) % mod )
inv[0]=0

_fft_mod = 998244353
_fft_imag = 911660635
_fft_iimag = 86583718
_fft_rate2 = (911660635, 509520358, 369330050, 332049552, 983190778, 123842337, 238493703, 975955924, 603855026, 856644456, 131300601,
              842657263, 730768835, 942482514, 806263778, 151565301, 510815449, 503497456, 743006876, 741047443, 56250497, 867605899)
_fft_irate2 = (86583718, 372528824, 373294451, 645684063, 112220581, 692852209, 155456985, 797128860, 90816748, 860285882, 927414960,
               354738543, 109331171, 293255632, 535113200, 308540755, 121186627, 608385704, 438932459, 359477183, 824071951, 103369235)
_fft_rate3 = (372528824, 337190230, 454590761, 816400692, 578227951, 180142363, 83780245, 6597683, 70046822, 623238099,
              183021267, 402682409, 631680428, 344509872, 689220186, 365017329, 774342554, 729444058, 102986190, 128751033, 395565204)
_fft_irate3 = (509520358, 929031873, 170256584, 839780419, 282974284, 395914482, 444904435, 72135471, 638914820, 66769500,
               771127074, 985925487, 262319669, 262341272, 625870173, 768022760, 859816005, 914661783, 430819711, 272774365, 530924681)
 
 
def _butterfly(a):
    n = len(a)
    h = (n - 1).bit_length()
    len_ = 0
    while len_ < h:
        if h - len_ == 1:
            p = 1 << (h - len_ - 1)
            rot = 1
            for s in range(1 << len_):
                offset = s << (h - len_)
                for i in range(p):
                    l = a[i + offset]
                    r = a[i + offset + p] * rot % _fft_mod
                    a[i + offset] = (l + r) % _fft_mod
                    a[i + offset + p] = (l - r) % _fft_mod
                if s + 1 != (1 << len_):
                    rot *= _fft_rate2[(~s & -~s).bit_length() - 1]
                    rot %= _fft_mod
            len_ += 1
        else:
            p = 1 << (h - len_ - 2)
            rot = 1
            for s in range(1 << len_):
                rot2 = rot * rot % _fft_mod
                rot3 = rot2 * rot % _fft_mod
                offset = s << (h - len_)
                for i in range(p):
                    a0 = a[i + offset]
                    a1 = a[i + offset + p] * rot
                    a2 = a[i + offset + p * 2] * rot2
                    a3 = a[i + offset + p * 3] * rot3
                    a1na3imag = (a1 - a3) % _fft_mod * _fft_imag
                    a[i + offset] = (a0 + a2 + a1 + a3) % _fft_mod
                    a[i + offset + p] = (a0 + a2 - a1 - a3) % _fft_mod
                    a[i + offset + p * 2] = (a0 - a2 + a1na3imag) % _fft_mod
                    a[i + offset + p * 3] = (a0 - a2 - a1na3imag) % _fft_mod
                if s + 1 != (1 << len_):
                    rot *= _fft_rate3[(~s & -~s).bit_length() - 1]
                    rot %= _fft_mod
            len_ += 2
 
 
def _butterfly_inv(a):
    n = len(a)
    h = (n - 1).bit_length()
    len_ = h
    while len_:
        if len_ == 1:
            p = 1 << (h - len_)
            irot = 1
            for s in range(1 << (len_ - 1)):
                offset = s << (h - len_ + 1)
                for i in range(p):
                    l = a[i + offset]
                    r = a[i + offset + p]
                    a[i + offset] = (l + r) % _fft_mod
                    a[i + offset + p] = (l - r) * irot % _fft_mod
                if s + 1 != (1 << (len_ - 1)):
                    irot *= _fft_irate2[(~s & -~s).bit_length() - 1]
                    irot %= _fft_mod
            len_ -= 1
        else:
            p = 1 << (h - len_)
            irot = 1
            for s in range(1 << (len_ - 2)):
                irot2 = irot * irot % _fft_mod
                irot3 = irot2 * irot % _fft_mod
                offset = s << (h - len_ + 2)
                for i in range(p):
                    a0 = a[i + offset]
                    a1 = a[i + offset + p]
                    a2 = a[i + offset + p * 2]
                    a3 = a[i + offset + p * 3]
                    a2na3iimag = (a2 - a3) * _fft_iimag % _fft_mod
                    a[i + offset] = (a0 + a1 + a2 + a3) % _fft_mod
                    a[i + offset + p] = (a0 - a1 +
                                         a2na3iimag) * irot % _fft_mod
                    a[i + offset + p * 2] = (a0 + a1 -
                                             a2 - a3) * irot2 % _fft_mod
                    a[i + offset + p * 3] = (a0 - a1 -
                                             a2na3iimag) * irot3 % _fft_mod
                if s + 1 != (1 << (len_ - 1)):
                    irot *= _fft_irate3[(~s & -~s).bit_length() - 1]
                    irot %= _fft_mod
            len_ -= 2
 
 
def _convolution_naive(a, b):
    n = len(a)
    m = len(b)
    ans = [0] * (n + m - 1)
    if n < m:
        for j in range(m):
            for i in range(n):
                ans[i + j] = (ans[i + j] + a[i] * b[j]) % _fft_mod
    else:
        for i in range(n):
            for j in range(m):
                ans[i + j] = (ans[i + j] + a[i] * b[j]) % _fft_mod
    return ans
 
 
def _convolution_fft(a, b):
    a = a.copy()
    b = b.copy()
    n = len(a)
    m = len(b)
    z = 1 << (n + m - 2).bit_length()
    a += [0] * (z - n)
    _butterfly(a)
    b += [0] * (z - m)
    _butterfly(b)
    for i in range(z):
        a[i] = a[i] * b[i] % _fft_mod
    _butterfly_inv(a)
    a = a[:n + m - 1]
    iz = pow(z, _fft_mod - 2, _fft_mod)
    for i in range(n + m - 1):
        a[i] = a[i] * iz % _fft_mod
    return a
 
 
def _convolution_square(a):
    a = a.copy()
    n = len(a)
    z = 1 << (2 * n - 2).bit_length()
    a += [0] * (z - n)
    _butterfly(a)
    for i in range(z):
        a[i] = a[i] * a[i] % _fft_mod
    _butterfly_inv(a)
    a = a[:2 * n - 1]
    iz = pow(z, _fft_mod - 2, _fft_mod)
    for i in range(2 * n - 1):
        a[i] = a[i] * iz % _fft_mod
    return a
 
 
def convolution(a, b):
    """It calculates (+, x) convolution in mod 998244353. 
    Given two arrays a[0], a[1], ..., a[n - 1] and b[0], b[1], ..., b[m - 1], 
    it calculates the array c of length n + m - 1, defined by
 
    >   c[i] = sum(a[j] * b[i - j] for j in range(i + 1)) % 998244353.
 
    It returns an empty list if at least one of a and b are empty.
 
    Constraints
    -----------
 
    >   len(a) + len(b) <= 8388609
 
    Complexity
    ----------
 
    >   O(n log n), where n = len(a) + len(b).
    """
    n = len(a)
    m = len(b)
    if n == 0 or m == 0:
        return []
    if min(n, m) <= 0:
        return _convolution_naive(a, b)
    if a is b:
        return _convolution_square(a)
    return _convolution_fft(a, b)

def find_primitive_root(P):
    while True:
        g = random.randint(1,P-1)
        tmp = [0] * P
        tmp[1] = True
        a = 1
        while tmp[(a*g) % P] == 0:
            tmp[(a * g) % P] = 1
            a = (a * g) % P
        if sum(tmp) == P-1:
            return g

N,P = mi()

fact = [1] * (P)
inv = [1] * P
inv[0] = 0
for i in range(2,P):
    fact[i] = i * fact[i-1] % P
    inv[i] = (-inv[P % i]) * (P//i) % P
ifact = [inv[fact[i]] for i in range(P)]

g = find_primitive_root(P)
log_g = [-1] * P
for i in range(P-1):
    log_g[pow(g,i,P)] = i

f = [0] * (P-1)
f[0] = 1
while N:
    n = N % P
    tmp = [0] * (P-1)
    for k in range(n+1):
        val = fact[n] * (ifact[k] * ifact[n-k] % P) % P
        tmp[log_g[val]] += 1
    f = convolution(f,tmp)
    for i in range(P-1,len(f)):
        f[i % (P-1)] += f[i]
        f[i % (P-1)] %= mod
    f = f[:P-1]
    N //= P

res = 0
for i in range(P-1):
    res += pow(g,i,P) * f[i] % mod
    res %= mod

print(res)

yukicoder

結果

テストケース

ソースコード