結果

問題 No.2959 Dolls' Tea Party
ユーザー gew1fw
提出日時 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
権限があれば一括ダウンロードができます

ソースコード

diff #

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()
0