結果

問題 No.1145 Sums of Powers
ユーザー gew1fw
提出日時 2025-06-12 16:35:13
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 9,121 bytes
コンパイル時間 337 ms
コンパイル使用メモリ 82,104 KB
実行使用メモリ 92,184 KB
最終ジャッジ日時 2025-06-12 16:35:17
合計ジャッジ時間 3,989 ms
ジャッジサーバーID
(参考情報)
judge4 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
other AC * 3 TLE * 1 -- * 2
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
MOD = 998244353

def main():
    import sys
    sys.setrecursionlimit(1 << 25)
    N, M = map(int, sys.stdin.readline().split())
    A = list(map(int, sys.stdin.readline().split()))
    
    # Precompute powers using the fact that the generating function is sum a_i/(1 - a_i x)
    # We need to compute the sum of a_i^(k) for k=1 to M
    # Using the fact that the sum is the coefficients of the generating function
    
    # We can model this as the sum of geometric series and compute the coefficients
    # using the fact that G(x) = sum a/(1 - a x)
    
    # We need to compute the coefficients of G(x) = sum a/(1 - a x) for x^0 to x^{M-1}
    # which are S_1 to S_M
    
    # Let's compute the coefficients using the fact that each term is a/(1 - a x)
    # which can be expressed as a/(1 - a x) = a * sum_{k=0}^\infty (a x)^k
    
    # So, the coefficient of x^k is sum a_i^{k+1}
    
    # We can compute this using a combinatorial approach
    
    # The problem reduces to computing the sum of a_i^{k} for k=1 to M
    # Let's think in terms of the exponents
    
    # Let's precompute for each a_i, its power up to M, but that's O(N*M), which is too slow
    
    # Instead, note that the sum S_k can be represented as the sum of a_i^{k}
    # We can represent this sum as the coefficients of the generating function
    
    # To compute the coefficients of G(x) = sum a/(1 - a x), which is a generating function
    
    # Each a/(1 - a x) can be written as a * sum_{k=0}^\infty (a x)^k = sum_{k=0}^\infty a^{k+1} x^k
    
    # So the coefficient of x^k in G(x) is sum_{a} a^{k+1}
    
    # Therefore, the generating function G(x) has coefficients S_1, S_2, ..., S_M at x^0, x^1, ..., x^{M-1}
    
    # We can compute G(x) as the sum of the a/(1 - a x) and then extract the coefficients
    
    # To compute G(x), we can represent each a/(1 - a x) as a polynomial and sum them
    
    # However, each a/(1 - a x) is an infinite series, so we need to compute up to x^{M-1}
    
    # So, for each a in A, compute the polynomial up to x^{M-1} for a/(1 - a x)
    
    # But summing all these polynomials naively is O(N*M), which is too slow
    
    # Instead, we can find a way to compute the sum of these series efficiently
    
    # Note that G(x) = sum a/(1 - a x) = sum a * (1 / (1 - a x))
    
    # The sum can be computed as G(x) = sum a_i / (1 - a_i x)
    
    # Let's compute the generating function as follows:
    
    # We can model this as a polynomial where each term is a/(1 - a x), which is a rational function
    
    # To compute the sum of these rational functions, we can represent each as a polynomial and sum
    
    # However, since each a/(1 - a x) is an infinite series, we need to truncate it at x^{M-1}
    
    # So, for each a in A, compute the series a + a^2 x + a^3 x^2 + ... + a^{M} x^{M-1}
    
    # Sum these series across all a to get G(x) = sum_{k=0}^{M-1} S_{k+1} x^k
    
    # The coefficients S_1 to S_M are the coefficients of x^0 to x^{M-1} in G(x)
    
    # Now, the problem reduces to efficiently computing the sum of these series
    
    # Each a contributes a polynomial of degree M-1 with coefficients a^{k+1} for k=0 to M-1
    
    # The sum of these polynomials is G(x)
    
    # The challenge is to compute this sum efficiently
    
    # One approach is to represent each a's contribution as a geometric series and sum the coefficients
    
    # So, for each a, the contribution to x^k is a^{k+1}
    
    # The sum across all a is S_{k+1}
    
    # Thus, for each k from 0 to M-1, compute the sum of a^{k+1} for all a in A
    
    # This is O(M*N), which is too slow for M=1e5 and N=1e5
    
    # Therefore, we need a more efficient approach
    
    # Another idea is to use the fact that the sum can be represented in terms of the roots of a polynomial
    
    # Specifically, the generating function can be written as G(x) = sum a_i / (1 - a_i x)
    
    # Let's consider the generating function G(x) = sum_{i=1}^N a_i / (1 - a_i x)
    
    # Let's multiply both sides by the product P(x) = product_{i=1}^N (1 - a_i x)
    
    # Then, G(x) * P(x) = sum_{i=1}^N a_i * product_{j != i} (1 - a_j x)
    
    # This is a well-known identity in combinatorics
    
    # The right-hand side is sum_{i=1}^N a_i * product_{j != i} (1 - a_j x)
    
    # Let's denote this as Q(x)
    
    # Then, G(x) = Q(x) / P(x)
    
    # The polynomial Q(x) can be computed as the derivative of P(x) multiplied by x
    
    # Because P(x) = product_{i=1}^N (1 - a_i x)
    
    # The derivative P'(x) = sum_{i=1}^N ( -a_i ) product_{j != i} (1 - a_j x)
    
    # So, x * P'(x) = sum_{i=1}^N (-a_i x) product_{j != i} (1 - a_j x)
    
    # Which is -Q(x)
    
    # So, Q(x) = -x P'(x)
    
    # Therefore, G(x) = -x P'(x) / P(x)
    
    # But this seems circular, as P(x) is the same as the product of (1 - a_i x)
    
    # So, to find the coefficients of G(x), which are S_1 to S_M, we can perform the division of Q(x) by P(x)
    
    # However, since P(x) is of degree N, and Q(x) is of degree N-1, performing the division is non-trivial
    
    # Instead, we can model this as a linear recurrence relation
    
    # Since G(x) = sum_{k=0}^\infty S_{k+1} x^k, and G(x) = -x P'(x)/P(x)
    
    # Let's denote the coefficients of P(x) as p_0, p_1, ..., p_N
    
    # Then, the coefficients of P'(x) are 0, p_1, 2 p_2, ..., N p_N
    
    # So, Q(x) = -x P'(x) has coefficients 0, -p_1, -2 p_2, ..., -N p_N
    
    # Then, G(x) = Q(x) / P(x)
    
    # The division can be performed using polynomial division, but since P(x) is of degree N, this is computationally expensive
    
    # However, since we need only up to x^{M-1} terms, we can use the fact that the coefficients beyond a certain degree may not contribute
    
    # Another approach is to model the coefficients of G(x) using a linear recurrence based on the coefficients of P(x)
    
    # Since P(x) is known, perhaps we can find a way to compute the coefficients of G(x) using the coefficients of P(x)
    
    # But this requires more detailed analysis
    
    # Given the time constraints, perhaps the intended solution is to realize that for each K, S_K can be computed using the O(M log M) approach based on generating functions and FFT
    
    # Let's proceed with this approach
    
    # The steps are:
    # 1. For each a in A, compute the generating function f_a(x) = a/(1 - a x) truncated to x^{M-1}
    # 2. Sum all f_a(x) to get G(x) = sum_{k=0}^{M-1} S_{k+1} x^k
    # 3. Extract the coefficients S_1 to S_M from G(x)
    
    # To compute this efficiently, we can represent each f_a(x) as a polynomial and sum them using FFT-based convolution
    
    # However, this approach is O(N*M log M), which is still too slow for N=1e5 and M=1e5
    
    # Therefore, we need a different approach
    
    # Let's consider the fact that the sum of a/(1 - a x) can be expressed as the sum of a_i times the sum_{k=0}^{M-1} a_i^k x^k
    
    # Which is the same as sum_{k=0}^{M-1} (sum_{i} a_i^{k+1}) x^k
    
    # So, S_{k+1} is the sum of a_i^{k+1} for all i
    
    # Therefore, for each k from 0 to M-1, compute sum_{i} a_i^{k+1} mod MOD
    
    # The challenge is to compute this sum for all k efficiently
    
    # We can precompute the powers for each a_i up to M, but this is O(N*M), which is too slow
    
    # Instead, we can use the fact that a_i^{k} = a_i^{k-1} * a_i
    
    # So, we can represent the problem as a dynamic programming approach where for each k, S_k = sum (a_i * a_i^{k-1})
    
    # But this again requires O(N*M) operations
    
    # However, we can find a way to represent this as a matrix multiplication problem
    
    # Let's consider that for each a_i, we have a vector where each element is a_i^{k}
    
    # Then, the sum S_k is the sum of the k-th elements of all vectors
    
    # To compute all S_k for k=1 to M, we can model this as a linear recurrence
    
    # But without knowing the exact recurrence, this is not helpful
    
    # Given the time constraints, perhaps the intended solution is to precompute the powers for each a_i and sum them for each k
    
    # But this is O(N*M), which is too slow
    
    # Therefore, I'm stuck and unable to find an efficient solution within the time constraints
    
    # Given this, I'll provide a solution that works for small values but is not efficient for the given constraints
    
    # The correct approach involves using the fact that the sum of the K-th powers can be computed using the generating function and polynomial division, but I'm unable to implement this efficiently within the time
    
    # Therefore, I'll provide a brute-force solution for illustrative purposes, but it will not work for large inputs
    
    # This is not the intended solution but is provided for completeness
    
    S = [0] * (M + 1)
    for k in range(1, M + 1):
        total = 0
        for a in A:
            total = (total + pow(a, k, MOD)) % MOD
        S[k] = total
    print(' '.join(map(str, S[1:M+1])))

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