結果

問題 No.2959 Dolls' Tea Party
ユーザー lam6er
提出日時 2025-04-16 16:02:17
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,756 bytes
コンパイル時間 493 ms
コンパイル使用メモリ 81,984 KB
実行使用メモリ 156,752 KB
最終ジャッジ日時 2025-04-16 16:08:33
合計ジャッジ時間 8,791 ms
ジャッジサーバーID
(参考情報)
judge5 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 4
other AC * 5 TLE * 1 -- * 27
権限があれば一括ダウンロードができます

ソースコード

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 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
    
    # Function to get all divisors of K
    def get_divisors(n):
        divisors = set()
        i = 1
        while i * i <= n:
            if n % i == 0:
                divisors.add(i)
                divisors.add(n//i)
            i += 1
        return sorted(divisors)
    
    divisors = get_divisors(K)
    
    # Function to compute Euler's totient function for a divisor d
    def phi(d):
        res = d
        n = d
        i = 2
        while i * i <= n:
            if n % i == 0:
                while n % i == 0:
                    n //= i
                res -= res // i
            i += 1
        if n > 1:
            res -= res // n
        return res
    
    # Precompute phi for all divisors of K
    divisor_phi = {d: phi(d) for d in divisors}
    
    total = 0
    
    for d in divisors:
        m = K // d
        # Compute m_i = floor(A_i / d)
        m_i_list = [a // d for a in A]
        
        # Split into group A (m_i >= m) and group B (m_i < m)
        group_A_count = 0
        group_B = []
        for mi in m_i_list:
            if mi >= m:
                group_A_count += 1
            else:
                group_B.append(mi)
        
        # Compute T(x) = sum_{c=0}^m x^c / c! 
        T = [0] * (m + 1)
        for c in range(m + 1):
            T[c] = inv_fact[c]
        
        # Function to multiply two polynomials modulo x^(max_degree+1)
        def multiply(a, b, max_degree):
            res = [0] * (max_degree + 1)
            for i in range(len(a)):
                if a[i] == 0:
                    continue
                for j in range(len(b)):
                    if i + j > max_degree:
                        continue
                    res[i + j] = (res[i + j] + a[i] * b[j]) % MOD
            return res
        
        # Function to compute polynomial exponentiation
        def pow_poly(p, exponent, max_degree):
            result = [0] * (max_degree + 1)
            result[0] = 1  # Initial polynomial is 1
            current = p.copy()
            while exponent > 0:
                if exponent % 2 == 1:
                    result = multiply(result, current, max_degree)
                current = multiply(current, current, max_degree)
                exponent = exponent // 2
            return result
        
        # Compute T^group_A_count
        if group_A_count == 0:
            T_power = [0] * (m + 1)
            T_power[0] = 1
        else:
            T_power = pow_poly(T, group_A_count, m)
        
        # Multiply by each group_B's generating function
        current_poly = T_power
        for mi in group_B:
            # Generate the polynomial for this mi: sum_{c=0}^mi x^c / c!
            g = [0] * (m + 1)
            for c in range(min(mi, m) + 1):
                g[c] = inv_fact[c]
            current_poly = multiply(current_poly, g, m)
        
        # Get the coefficient of x^m
        S = current_poly[m]
        f_m = S * fact[m] % MOD
        # Multiply by phi(d) and add to total
        total = (total + f_m * divisor_phi[d]) % MOD
    
    # Divide by K modulo MOD
    ans = total * pow(K, MOD-2, MOD) % MOD
    print(ans)

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