結果

問題 No.1357 Nada junior high school entrance examination 3rd day
ユーザー qwewe
提出日時 2025-05-14 13:14:16
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,043 ms / 2,000 ms
コード長 7,897 bytes
コンパイル時間 307 ms
コンパイル使用メモリ 82,288 KB
実行使用メモリ 105,680 KB
最終ジャッジ日時 2025-05-14 13:15:44
合計ジャッジ時間 15,942 ms
ジャッジサーバーID
(参考情報)
judge5 / judge4
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 1
other AC * 21
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys

# Set a higher recursion depth, just in case. With iterative NTT, it might not be strictly necessary.
# sys.setrecursionlimit(200000) 

def solve():
    K = int(sys.stdin.readline())
    MOD = 998244353

    # The problem asks for coefficients c_0, c_1, ..., c_{2K}.
    # The maximum index is 2K. The total number of coefficients is 2K+1.
    MAX_N = 2 * K 
    
    # Precomputation requires factorials up to 2K+1.
    # Because the polynomial A(t) has coefficients a_n = 1/(n+1)!, 
    # for the highest degree term needed n=2K, we need a_{2K} = 1/(2K+1)!.
    # So we need factorials and inverse factorials up to 2K+1.
    MAX_FACT = MAX_N + 1 
    
    # Handle K=0 case if it were possible. The problem states K >= 1.
    # If K=0, MAX_N = 0, MAX_FACT = 1. The sum is empty, result is 0. Output should be just c_0=0.
    if K == 0:
         print("0") 
         return

    # Precompute factorials and their modular inverses.
    # Size of arrays needs to be MAX_FACT + 1 to store indices 0..MAX_FACT.
    fact = [1] * (MAX_FACT + 1)
    inv_fact = [1] * (MAX_FACT + 1)
    
    for i in range(1, MAX_FACT + 1):
        fact[i] = (fact[i-1] * i) % MOD
    
    # Compute modular inverse using Fermat's Little Theorem for the largest factorial needed.
    # (MOD - 2) is the exponent for modular inverse because MOD is prime.
    inv_fact[MAX_FACT] = pow(fact[MAX_FACT], MOD - 2, MOD)
    
    # Compute other inverse factorials iteratively using inv_fact[i] = inv_fact[i+1] * (i+1)
    for i in range(MAX_FACT - 1, -1, -1):
        inv_fact[i] = (inv_fact[i+1] * (i+1)) % MOD

    # Iterative Number Theoretic Transform (NTT)
    # This implementation uses Cooley-Tukey Radix-2 DIT FFT algorithm adapted for modular arithmetic.
    def ntt(a, inv):
        n = len(a)
        if n == 1:
            return

        # Bit-reversal permutation: arrange elements in order required for iterative FFT
        j = 0
        for i in range(1, n):
            rev = n >> 1
            while j >= rev:
                j -= rev
                rev >>= 1
            j += rev
            if i < j:
                a[i], a[j] = a[j], a[i]

        # Main FFT loop
        len_ = 2 # Current length of DFT sequences being combined
        while len_ <= n:
            # Root of unity: g=3 is primitive root for MOD 998244353
            # w_len is omega_len_ = g^((MOD-1)/len_) mod MOD
            w_len = pow(3, (MOD - 1) // len_, MOD)
            if inv:
                # For inverse NTT, use inverse root of unity
                w_len = pow(w_len, MOD - 2, MOD)
            
            half_len = len_ >> 1
            for i in range(0, n, len_): # Iterate through blocks of size len_
                w = 1 # Current power of root of unity: omega_len_^j
                for j in range(half_len): # Butterfly operations
                    u = a[i + j] # Even indexed element
                    v = (w * a[i + j + half_len]) % MOD # Odd indexed element multiplied by root power
                    a[i + j] = (u + v) % MOD
                    a[i + j + half_len] = (u - v + MOD) % MOD # Ensure result is non-negative
                    w = (w * w_len) % MOD # Update root power for next step
            len_ <<= 1 # Double the length for next stage

        if inv:
            # Scale by 1/n for inverse transform
            n_inv = pow(n, MOD - 2, MOD)
            for i in range(n):
                a[i] = (a[i] * n_inv) % MOD

    # Polynomial inverse computation using NTT (based on Newton-Raphson iteration)
    def poly_inv(A, deg):
        # Computes polynomial B such that A * B = 1 (mod x^deg)
        
        # Pad A to ensure it has length at least deg
        A_padded = A[:deg] + [0] * max(0, deg - len(A))
        
        B = [0] * deg # Result polynomial B
        if deg == 0: return B 
        
        # Base case: B = 1/A[0] (mod x^1)
        B[0] = pow(A_padded[0], MOD - 2, MOD)
        
        # Iteratively double the precision `k`
        k = 1
        while k < deg:
            k *= 2
            
            # Determine required NTT length: smallest power of 2 >= 2k
            # This is needed because multiplication of two polynomials of degree k-1 results in degree 2k-2.
            ntt_len = 1 << (2*k -1).bit_length() if k > 0 else 1

            # Prepare A_k and B_k for NTT: Take required parts and pad to ntt_len
            # We only need A up to degree k-1 for computation mod x^k
            A_k = A_padded[:min(k, deg)] + [0] * (ntt_len - min(k, deg))
            # Use current approximation B up to degree k-1
            B_k = B[:min(k, deg)] + [0] * (ntt_len - min(k, deg)) 

            ntt(A_k, False)
            ntt(B_k, False)

            # Compute C_ntt = NTT(A * B^2) efficiently in frequency domain
            # C_ntt[i] = A_k[i] * B_k[i] * B_k[i]
            C_ntt = [0] * ntt_len
            for i in range(ntt_len):
                C_ntt[i] = (A_k[i] * B_k[i]) % MOD 
                C_ntt[i] = (C_ntt[i] * B_k[i]) % MOD 

            ntt(C_ntt, True) # Inverse NTT to get coefficients of A*B^2

            # Update B using the Newton-Raphson formula: B_new = 2B - A*B^2 (mod x^k)
            # We need to update B up to degree k-1. Copy the relevant part of B first.
            B_new_k = B[:k] 
            
            for i in range(min(k, deg)): # Compute updates up to degree min(k, deg) - 1
                 term_2B = (2 * B[i]) % MOD
                 # C_ntt holds coefficients of A*B^2. Indices correspond to powers of x.
                 B_new_k[i] = (term_2B - C_ntt[i] + MOD) % MOD 
            
            # Update B for next iteration. Store the result padded to degree `deg`.
            B = B_new_k + [0] * max(0, deg - k)

        return B[:deg] # Return the inverse polynomial up to degree deg-1

    # Step 1: Define polynomial A(t) coefficients a_n = 1 / (n+1)!
    # We need polynomial of degree 2K, so length is 2K+1.
    N_poly = MAX_N + 1 
    
    A = [0] * N_poly
    for n in range(N_poly):
         # A[n] is the coefficient of t^n
         A[n] = inv_fact[n+1] # Access precomputed inverse factorials: 1/(n+1)!

    # Step 2: Compute B(t) = A(t)^(-1) mod t^(N_poly)
    # The coefficients b_n of B(t) are B_n / n! (where B_n are Bernoulli numbers with B1=-1/2 convention)
    inv_deg = N_poly # Compute inverse modulo t^(2K+1)
    B_coeffs = poly_inv(A, inv_deg) 

    # Step 3: Compute the final coefficients c_i based on the derived formula
    # c_{2a} = (2a-1) * (-1)^(a-1) * (B_{2a} / (2a)!) * 2^(2a-1)
    C = [0] * (MAX_N + 1) # Result array C[0]...C[2K]
    
    # Precompute powers of 2 modulo MOD
    pow2 = [1] * (MAX_N + 1)
    for i in range(1, MAX_N + 1):
        pow2[i] = (pow2[i-1] * 2) % MOD

    # Calculate coefficients c_{2a} for a = 1 to K
    for a in range(1, K + 1):
        idx = 2*a # Index in C array corresponds to power of pi (pi^idx)
        
        # Safety check: Index should be within computed range
        if idx >= N_poly: continue 
        
        term1 = (2*a - 1) # The factor (2a-1)
        
        term2 = 1 # The factor (-1)^(a-1)
        if (a - 1) % 2 == 1: # Check if a-1 is odd (i.e., a is even)
             term2 = MOD - 1 # Represents -1 mod P
        
        # b_{2a} = B_{2a} / (2a)! is the coefficient from poly inverse B_coeffs[idx]
        b_2a = B_coeffs[idx] 
        
        idx_pow2 = idx - 1 # Index for power of 2 is 2a-1
        # Safety check: index for power of 2 must be non-negative
        if idx_pow2 < 0 : continue 
        term4 = pow2[idx_pow2] # Factor 2^(2a-1)

        # Combine all terms modulo MOD
        c_2a = (term1 % MOD * term2) % MOD
        c_2a = (c_2a * b_2a) % MOD
        c_2a = (c_2a * term4) % MOD
        
        C[idx] = c_2a # Store the computed coefficient c_{2a} at index 2a

    # Step 4: Print the result coefficients C[0], C[1], ..., C[2K] separated by spaces
    print(*(C))

solve()
0