結果

問題 No.3045 反復重み付き累積和
ユーザー Nauclhlt🪷
提出日時 2025-02-01 16:05:15
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 2,171 bytes
コンパイル時間 548 ms
コンパイル使用メモリ 82,260 KB
実行使用メモリ 138,520 KB
最終ジャッジ日時 2025-02-01 16:09:08
合計ジャッジ時間 231,835 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1 WA * 1 TLE * 1
other WA * 4 TLE * 37
権限があれば一括ダウンロードができます

ソースコード

diff # 

MOD = 998244353

def weighted_cumulative_sum(A, k):
    """Calculates the weighted cumulative sum of A with weight k efficiently."""
    n = len(A)
    P = [0] * n
    for i in range(n):
        coeff = 1
        for j in range(i, -1, -1):
            P[i] = (P[i] + coeff * A[j]) % MOD
            coeff = (coeff * k) % MOD
    return P

def power_weighted_cumulative_sum(A, k, x):
    """Efficiently applies the weighted cumulative sum operation x times."""
    n = len(A)
    if x == 0:
        return A
    if x == 1:
        return weighted_cumulative_sum(A, k)

    # Precompute powers of k
    k_powers = [1] * n
    for i in range(1, n):
        k_powers[i] = (k_powers[i - 1] * k) % MOD

    # Apply cumulative sum iteratively using powers of k
    for _ in range(x):
        P = [0] * n
        for i in range(n):
            coeff = 1
            for j in range(i, -1, -1):
                P[i] = (P[i] + coeff * A[j]) % MOD
                coeff = (coeff * k_powers[i - j]) % MOD
        A = P
    return A

def process_queries(N, Q, A, queries):
    result = []
    for query in queries:
        if query[0] == 1:
            # Perform the "weighted cumulative sum" update on A with weight k, x times
            k, x = query[1], query[2]
            A = power_weighted_cumulative_sum(A, k, x)
        elif query[0] == 2:
            # Output the x-th element (1-indexed) modulo MOD
            x = query[1] - 1  # Convert to 0-based index
            result.append(A[x] % MOD)
    return result

# Input processing
import sys
input = sys.stdin.read

data = input().split()
N = int(data[0])
Q = int(data[1])
A = list(map(int, data[2:2 + N]))
queries = []
index = 2 + N
for _ in range(Q):
    query_type = int(data[index])
    if query_type == 1:
        # Read "1 k x" query
        k = int(data[index + 1])
        x = int(data[index + 2])
        queries.append([1, k, x])
        index += 3
    elif query_type == 2:
        # Read "2 x" query
        x = int(data[index + 1])
        queries.append([2, x])
        index += 2

# Process and output results
results = process_queries(N, Q, A, queries)
sys.stdout.write("\n".join(map(str, results)) + "\n")
0