結果

問題 No.2215 Slide Subset Sum
ユーザー lam6er
提出日時 2025-03-31 17:53:39
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,470 ms / 3,000 ms
コード長 2,171 bytes
コンパイル時間 288 ms
コンパイル使用メモリ 82,440 KB
実行使用メモリ 268,772 KB
最終ジャッジ日時 2025-03-31 17:55:14
合計ジャッジ時間 33,970 ms
ジャッジサーバーID
(参考情報)
judge4 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 45
権限があれば一括ダウンロードができます

ソースコード

diff #

MOD = 998244353

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    ptr = 0
    N = int(data[ptr])
    ptr += 1
    M = int(data[ptr])
    ptr += 1
    K = int(data[ptr])
    ptr += 1
    A = list(map(int, data[ptr:ptr + N]))
    ptr += N

    front_stack = []  # Each element is (x, previous_dp)
    back_stack = []

    # Initialize DP arrays
    front_dp = [1] + [0] * (K - 1)
    back_dp = [1] + [0] * (K - 1)

    # Process the initial window
    for i in range(M):
        x = A[i]
        back_stack.append(x)
        new_back_dp = [0] * K
        for r in range(K):
            new_back_dp[r] = (back_dp[r] + back_dp[(r - x) % K]) % MOD
        back_dp = new_back_dp

    def compute_convolution():
        total = 0
        for r in range(K):
            s = (K - r) % K
            total = (total + front_dp[r] * back_dp[s]) % MOD
        return (total - 1) % MOD

    answers = []
    answers.append(compute_convolution())

    # Slide the window
    for i in range(N - M):
        # Remove the leftmost element
        if not front_stack:
            # Transfer all elements from back to front
            while back_stack:
                x = back_stack.pop()
                prev_front_dp = front_dp.copy()
                new_front_dp = [0] * K
                for r in range(K):
                    new_front_dp[r] = (prev_front_dp[r] + prev_front_dp[(r - x) % K]) % MOD
                front_stack.append((x, prev_front_dp))
                front_dp = new_front_dp
            # Reset back_dp after transfer
            back_dp = [1] + [0] * (K - 1)
        
        # Pop from front_stack
        x_removed, prev_front = front_stack.pop()
        front_dp = prev_front

        # Add the new element to back_stack
        new_x = A[i + M]
        back_stack.append(new_x)
        new_back = [0] * K
        for r in range(K):
            new_back[r] = (back_dp[r] + back_dp[(r - new_x) % K]) % MOD
        back_dp = new_back

        # Compute and store the answer
        answers.append(compute_convolution())

    # Print all answers
    print('\n'.join(map(str, answers)))

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