結果
問題 |
No.2959 Dolls' Tea Party
|
ユーザー |
![]() |
提出日時 | 2025-06-12 20:43:46 |
言語 | PyPy3 (7.3.15) |
結果 |
TLE
|
実行時間 | - |
コード長 | 2,951 bytes |
コンパイル時間 | 166 ms |
コンパイル使用メモリ | 82,540 KB |
実行使用メモリ | 147,692 KB |
最終ジャッジ日時 | 2025-06-12 20:44:07 |
合計ジャッジ時間 | 5,345 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 4 |
other | AC * 2 TLE * 1 -- * 30 |
ソースコード
import sys from math import isqrt MOD = 998244353 def main(): N, K = map(int, sys.stdin.readline().split()) A = list(map(int, sys.stdin.readline().split())) # Precompute factorials and inverse factorials modulo MOD up to K max_fact = K 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 # Precompute divisors of K def get_divisors(n): divisors = set() for i in range(1, isqrt(n)+1): if n % i ==0: divisors.add(i) divisors.add(n//i) return sorted(divisors) divisors = get_divisors(K) # Precompute phi for each s (divisors) def compute_phi(s): if s ==0: return 0 res = s i =2 while i*i <= s: if s%i ==0: res -= res //i while s%i ==0: s //=i i +=1 if s>1: res -= res //s return res phi = {} for s in divisors: phi[s] = compute_phi(s) total =0 for s in divisors: d = K // s # Compute the maximum x_i for each i mx = [] valid = True for a in A: if a <0: valid = False break mx_i = a // s if mx_i > d: mx_i = d mx.append(mx_i) if not valid: continue # Now, compute the generating function product # Initialize dp: dp[x] = coefficient of t^x dp = [0]*(d+1) dp[0] =1 for i in range(N): a = A[i] mx_i = mx[i] # Compute the generating function for this i: sum_{x=0}^{mx_i} t^x /x! # Since x can be up to min(mx_i, d), we cap it at d current_mx = min(mx_i, d) # For each x from current_mx down to 0: # Create a temporary array to store new dp values new_dp = [0]*(d+1) for x in range(d+1): if dp[x] ==0: continue # Add x + y contributions for y in range(0, current_mx +1): if x + y >d: break term = dp[x] * inv_fact[y] % MOD new_dp[x + y] = (new_dp[x + y] + term) % MOD dp = new_dp # Extract the coefficient for t^d coeff = dp[d] sum_multinom = coeff * fact[d] % MOD # Multiply by phi(s) contribution = sum_multinom * phi[s] % MOD total = (total + contribution) % MOD # Compute the inverse of K inv_K = pow(K, MOD-2, MOD) ans = total * inv_K % MOD print(ans) if __name__ == "__main__": main()