結果

問題 No.2005 Sum of Power Sums
コンテスト
ユーザー gew1fw
提出日時 2025-06-12 14:02:27
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 543 ms / 2,000 ms
コード長 1,857 bytes
コンパイル時間 537 ms
コンパイル使用メモリ 82,508 KB
実行使用メモリ 290,372 KB
最終ジャッジ日時 2025-06-12 14:03:17
合計ジャッジ時間 12,292 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 18
権限があれば一括ダウンロードができます

ソースコード

diff # 

MOD = 998244353

def main():
    import sys
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    M = int(input[ptr])
    ptr += 1
    K_list = list(map(int, input[ptr:ptr+N]))
    ptr += N

    max_k = 5000

    # Precompute Stirling numbers of the second kind
    stirling = [[0] * (max_k + 1) for _ in range(max_k + 1)]
    stirling[0][0] = 1
    for k in range(1, max_k + 1):
        for m in range(1, k + 1):
            stirling[k][m] = (m * stirling[k-1][m] + stirling[k-1][m-1]) % MOD

    # Precompute factorials and inverse factorials up to N + 5000
    max_fact = N + 5000
    fact = [1] * (max_fact + 1)
    for i in range(1, max_fact + 1):
        fact[i] = fact[i-1] * i % MOD

    inv_fact = [1] * (max_fact + 1)
    inv_fact[max_fact] = pow(fact[max_fact], MOD-2, MOD)
    for i in range(max_fact - 1, -1, -1):
        inv_fact[i] = inv_fact[i+1] * (i+1) % MOD

    # Compute M_mod
    M_mod = M % MOD

    # Precompute P(t) for t up to N + 5000
    max_t = N + 5000
    P = [1] * (max_t + 1)
    for t in range(1, max_t + 1):
        term = (M_mod + N - (t - 1)) % MOD
        P[t] = (P[t-1] * term) % MOD

    # Precompute C(t) for t up to N + 5000
    C = [0] * (max_t + 1)
    for t in range(0, max_t + 1):
        C[t] = P[t] * inv_fact[t] % MOD

    # Precompute pre_sum for all K up to 5000
    pre_sum = [0] * (max_k + 1)
    for K in range(0, max_k + 1):
        s = 0
        for m in range(0, K + 1):
            s_km = stirling[K][m]
            if s_km == 0:
                continue
            term = s_km * fact[m] % MOD
            t = N + m
            c = C[t]
            s = (s + term * c) % MOD
        pre_sum[K] = s

    # Calculate the answer
    answer = 0
    for K in K_list:
        answer = (answer + pre_sum[K]) % MOD
    print(answer)

if __name__ == "__main__":
    main()
0