結果
| 問題 |
No.2005 Sum of Power Sums
|
| ユーザー |
 qwewe
|
| 提出日時 | 2025-05-14 13:09:14 |
| 言語 | PyPy3 (7.3.15) |
| 結果 |
AC
|
| 実行時間 | 1,888 ms / 2,000 ms |
| コード長 | 6,927 bytes |
| コンパイル時間 | 165 ms |
| コンパイル使用メモリ | 82,592 KB |
| 実行使用メモリ | 278,076 KB |
| 最終ジャッジ日時 | 2025-05-14 13:10:21 |
| 合計ジャッジ時間 | 10,734 ms |
|
ジャッジサーバーID （参考情報） |
judge3 / judge4 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 3 |
| other | AC * 18 |
ソースコード
import sys
# Function to compute the solution
def solve():
# Read N and M from input
# N: length of the sequence A
# M: upper bound for the sum of elements in A
N, M = map(int, sys.stdin.readline().split())
# Read K values as a list
# K_i: the power to which A_i is raised
K = list(map(int, sys.stdin.readline().split()))
# Define the modulus
MOD = 998244353
# Find the maximum value in K. If K is empty (N=0), K_max defaults to 0.
# Problem constraints state N >= 1, so K is never empty.
K_max = 0
if N > 0: # Check N > 0 to handle potential empty K list case, although N >= 1.
K_max = max(K)
# The problem constraints state N >= 1, so N=0 is not possible.
# If N was 0, the only sequence is empty, sum is 0. Total sum is 0.
# Precompute factorials and inverse factorials modulo MOD.
# We need factorials up to N + K_max for calculating binomial coefficients.
MAX_FACT = N + K_max
# Initialize arrays for factorials and inverse factorials.
fact = [1] * (MAX_FACT + 1)
invFact = [1] * (MAX_FACT + 1)
# Compute factorials: fact[i] = i! mod MOD
for i in range(1, MAX_FACT + 1):
fact[i] = (fact[i-1] * i) % MOD
# Compute inverse factorial of MAX_FACT using Fermat's Little Theorem:
# a^(MOD-2) % MOD is the modular multiplicative inverse of a modulo MOD (if MOD is prime).
invFact[MAX_FACT] = pow(fact[MAX_FACT], MOD - 2, MOD)
# Compute inverse factorials downwards using the relation invFact[i] = invFact[i+1] * (i+1) mod MOD
for i in range(MAX_FACT - 1, -1, -1): # Iterate from MAX_FACT-1 down to 0
invFact[i] = (invFact[i+1] * (i+1)) % MOD
# Precompute Stirling numbers of the second kind S(k, j) using dynamic programming.
# S(k, j) is the number of ways to partition a set of k objects into j non-empty subsets.
# Stirling[k][j] stores S(k, j) mod MOD.
Stirling = [[0] * (K_max + 1) for _ in range(K_max + 1)]
# Base case: S(0, 0) = 1 (partitioning an empty set into 0 non-empty subsets)
if K_max >= 0: # Check needed if K_max could be negative (not possible here)
Stirling[0][0] = 1
# Compute S(k, j) using the recurrence relation: S(k, j) = S(k-1, j-1) + j * S(k-1, j)
for k in range(1, K_max + 1):
# Base case S(k, 0) = 0 for k >= 1 (cannot partition k items into 0 non-empty subsets)
Stirling[k][0] = 0
# Compute S(k, j) for j from 1 to k
for j in range(1, k + 1):
term1 = Stirling[k-1][j-1] # Represents partitioning k-1 items into j-1 subsets, then adding k-th item as new subset
term2 = (j * Stirling[k-1][j]) % MOD # Represents partitioning k-1 items into j subsets, then adding k-th item into one of the j subsets
Stirling[k][j] = (term1 + term2) % MOD
# The core idea is to compute the sum using the formula derived via Stirling numbers:
# Total Sum = sum_{j=0}^{K_max} [ Binomial(M+N, N+j) * j! * (sum_{i=1..N, K_i >= j} S(K_i, j)) ]
# Let C_j = Binomial(M+N, N+j) * j!
# Let S_j = sum_{i=1..N, K_i >= j} S(K_i, j)
# Total Sum = sum_{j=0}^{K_max} C_j * S_j mod P
# Precompute products needed for binomial coefficients with large upper term M+N.
# We need P_{N+j} = (M+N) * (M+N-1) * ... * (M+N - (N+j-1)) mod P
# P_{N+j} is the falling factorial (M+N)^{\underline{N+j}} mod P.
# Calculate M modulo MOD once. Python handles large integers M automatically.
M_mod = M % MOD
# Calculate P_N = product_{t=0}^{N-1} (M+N-t) mod P
P_N = 1
for t in range(N):
# Calculate term (M+N-t) mod P
term = (M + N - t) % MOD
# Python's % operator result is in [-MOD+1, MOD-1]. Add MOD to ensure non-negative before modulo if needed.
# (term % MOD + MOD) % MOD is the standard way to ensure positive result.
# But direct Python % works okay too: (-5 % 3 == 1)
P_N = (P_N * term) % MOD
# Calculate P_{N+j} values iteratively for j = 1..K_max using the relation:
# P_{N+j} = P_{N+j-1} * (M - j + 1) mod P
# Store P_{N+j} values in P_k_vals array. P_k_vals[j] stores P_{N+j}.
P_k_vals = [0] * (K_max + 1)
if K_max >= 0:
# Base case P_N corresponds to j=0
P_k_vals[0] = P_N
current_prod = P_N
for j in range(1, K_max + 1):
# Term to multiply is (M - (j-1)) = (M - j + 1)
term_to_mult = (M_mod - j + 1) % MOD
# Ensure term_to_mult is non-negative [0, MOD-1] range
term_to_mult = (term_to_mult + MOD) % MOD
current_prod = (current_prod * term_to_mult) % MOD
P_k_vals[j] = current_prod
# Compute C_j = Binomial(M+N, N+j) * j! mod P
C_j = [0] * (K_max + 1)
for j in range(K_max + 1):
# Lower index k of binomial coefficient Binomial(n, k)
idx_binom_lower = N + j
# Check boundary conditions for binomial coefficient Binomial(n, k)
# Need k >= 0 and n >= k.
# k = N+j. Since N>=1, j>=0, k>=1 is always true.
# n = M+N. Check if k > n, which is equivalent to N+j > M+N, or j > M.
if M < j:
C_j[j] = 0 # Binomial coefficient is 0 if k > n
continue
# Calculate Binomial(M+N, N+j) = P_{N+j} / (N+j)! = P_{N+j} * invFact[N+j]
# P_{N+j} = (M+N)^{\underline{N+j}} mod P is already computed as P_k_vals[j]
numerator = P_k_vals[j]
# Inverse factorial for the denominator
denominator_inv = invFact[idx_binom_lower]
# Compute the binomial coefficient value modulo MOD
binom_val = (numerator * denominator_inv) % MOD
# C_j = Binomial(M+N, N+j) * j!
C_j[j] = (binom_val * fact[j]) % MOD
# Compute counts N_k = number of times value k appears in the input list K
N_k_counts = [0] * (K_max + 1)
for k_val in K:
# Check constraints: K_i >= 1. Validate k_val is within expected range.
if 1 <= k_val <= K_max:
N_k_counts[k_val] += 1
# Compute S_j = sum_{k=j}^{K_max} N_k * S(k, j) mod P
# This sums the Stirling numbers S(k,j) weighted by the count of K_i=k, for all k >= j.
S_j = [0] * (K_max + 1)
for j in range(K_max + 1):
current_Sj = 0
# Sum over k from j up to K_max
for k in range(j, K_max + 1):
# Contribution from sequences where K_i = k
term = (N_k_counts[k] * Stirling[k][j]) % MOD
current_Sj = (current_Sj + term) % MOD
S_j[j] = current_Sj
# Compute the final answer = sum_{j=0}^{K_max} C_j * S_j mod P
total_sum = 0
for j in range(K_max + 1):
term = (C_j[j] * S_j[j]) % MOD
total_sum = (total_sum + term) % MOD
# Print the final result
print(total_sum)
# Call the solve function to execute the logic
solve()