結果

問題 No.2959 Dolls' Tea Party
ユーザー gew1fw
提出日時 2025-06-12 15:42:38
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,244 bytes
コンパイル時間 184 ms
コンパイル使用メモリ 82,772 KB
実行使用メモリ 130,376 KB
最終ジャッジ日時 2025-06-12 15:42:48
合計ジャッジ時間 6,181 ms
ジャッジサーバーID
(参考情報)
judge4 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 4
other AC * 2 TLE * 1 -- * 30
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
MOD = 998244353

def main():
    import sys
    sys.setrecursionlimit(1 << 25)
    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_m = K
    fact = [1] * (max_m + 1)
    for i in range(1, max_m + 1):
        fact[i] = fact[i-1] * i % MOD
    inv_fact = [1] * (max_m + 1)
    inv_fact[max_m] = pow(fact[max_m], MOD-2, MOD)
    for i in range(max_m-1, -1, -1):
        inv_fact[i] = inv_fact[i+1] * (i+1) % MOD

    # Function to find all divisors of K
    def get_divisors(k):
        divisors = set()
        for i in range(1, int(k**0.5) + 1):
            if k % i == 0:
                divisors.add(i)
                divisors.add(k//i)
        return sorted(divisors)
    divisors = get_divisors(K)

    # Precompute Euler's Totient function for each divisor
    from math import gcd
    def compute_phi(n):
        result = n
        i = 2
        while i*i <= n:
            if n % i == 0:
                while n % i == 0:
                    n //= i
                result -= result // i
            i += 1
        if n > 1:
            result -= result // n
        return result
    phi = {}
    for d in divisors:
        phi[d] = compute_phi(d)

    total = 0
    for d in divisors:
        m = K // d
        # Compute c_i = min(A_i // d, m)
        c = [min(a // d, m) for a in A]
        # Initialize dp
        dp = [0] * (m + 1)
        dp[0] = 1
        for ci in c:
            if ci == 0:
                continue
            # Create the generating function for this color
            s = [inv_fact[x] for x in range(ci + 1)]
            new_dp = [0] * (m + 1)
            for j in range(m + 1):
                for x in range(min(ci, j) + 1):
                    new_dp[j] = (new_dp[j] + dp[j - x] * s[x]) % MOD
            dp = new_dp
        # Compute the contribution for this d
        coeff = dp[m] % MOD
        ways = coeff * fact[m] % MOD
        ways = ways * phi[d] % MOD
        total = (total + ways) % MOD

    # The final answer is total / K mod MOD
    inv_K = pow(K, MOD-2, MOD)
    ans = total * inv_K % MOD
    print(ans)

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