Problem p02734 (python)

Fixed submission:
# https://atcoder.jp/contests/abc159/submissions/11139734

def main():
    MOD = 998244353

    N, S = map(int, input().split())
    *a, = map(int, input().split())

    dp = [0] * (S + 1)

    ret = 0
    for i, x in enumerate(a):
        dp[0] += 1
        if x > S: continue
        ret = (ret + dp[S - x] * (N - i)) % MOD
        for k in range(S - 1, x - 1, -1):
            dp[k] = (dp[k] + dp[k - x]) % MOD

    print(ret)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
def main():
    MOD = 998244353

    N, S = map(int, input().split())
    *A, = map(int, input().split())

    dp = [0] * (S + 1)

    ans = 0
    for x in A:
        ans += dp[S]
        if S - x >= 0:
            ans += dp[S - x]
        if S == x:
            ans += 1
        for s in range(S, x, -1):
            dp[s] = (dp[s] + dp[s - x]) % MOD
        dp[x] = (dp[x] + dp[0] + 1) % MOD
        dp[0] += 1
        ans %= MOD

    print(ans)


if __name__ == '__main__':
    main()

#             dp[i + 1][s][0] = (dp[i + 1][s][0] + dp[i][s][0]) % MOD
#             dp[i + 1][s][1] = (dp[i + 1][s][1]
#                                + dp[i][s][1] + (dp[i][s - x][1] if s - x >= 0 else 0)
#                                + dp[i][s][0] + (dp[i][s - x][0] if s - x >= 0 else 0)) % MOD
#             dp[i + 1][s][2] = (dp[i + 1][s][2]
#                                + dp[i][s][2] +
#                                dp[i][s][1] + (dp[i][s - x][1] if s - x >= 0 else 0)
#                                + dp[i][s][0] + (dp[i][s - x][0] if s - x >= 0 else 0)) % MOD

#     for i, x in enumerate(A):
#         dp[i + 1][0][0] = dp[i][0][0]
#         for s in range(S, -1, -1):
#             dp[i + 1][s][1] = (dp[i][s][1] + dp[i][s][0]) % MOD
#             dp[i + 1][s][2] = (dp[i][s][2] + dp[i][s][1] + dp[i][s][0]) % MOD
#
#             if s - x >= 0:
#                 v = dp[i][s - x][1] + dp[i][s - x][0]
#                 dp[i + 1][s][1] = (dp[i + 1][s][1] + v) % MOD
#                 dp[i + 1][s][2] = (dp[i + 1][s][2] + v) % MOD

#         dp[i + 1][0][0] = dp[i][0][0]
#         dp[i + 1][0][1] = (dp[i][0][1] + dp[i][0][0]) % MOD
#         dp[i + 1][0][2] = (dp[i][0][2] + dp[i][0][1] + dp[i][0][0]) % MOD
#         for s in range(S, 0, -1):
#             dp[i + 1][s][1] = dp[i][s][1] % MOD
#             dp[i + 1][s][2] = (dp[i][s][2] + dp[i][s][1]) % MOD
#             if s - x >= 0:
#                 v = dp[i][s - x][1] + dp[i][s - x][0]
#                 dp[i + 1][s][1] = (dp[i + 1][s][1] + v) % MOD
#                 dp[i + 1][s][2] = (dp[i + 1][s][2] + v) % MOD

#     for x in A:
#         for s in range(S, 0, -1):
#             dp[s][1] += dp[s][0]
#             if s - x >= 0:
#                 v = dp[s - x][0] + (1 if s - x == 0 else 0)
#                 dp[s][0] += v
#                 dp[s][1] += v
#             dp[s][0] %= MOD
#             dp[s][1] %= MOD
#         dp[0][0] += 1
#         dp[0][1] += (dp[0][0] + 1)

Buggy submission (verdict: Runtime Error):
# https://atcoder.jp/contests/abc159/submissions/11139734

def main():
    MOD = 998244353

    N, S = map(int, input().split())
    *a, = map(int, input().split())

    dp = [0] * (S + 1)

    ret = 0
    for i, x in enumerate(a):
        dp[0] += 1
        ret = (ret + dp[S - x] * (N - i)) % MOD
        for k in range(S - 1, x - 1, -1):
            dp[k] = (dp[k] + dp[k - x]) % MOD

    print(ret)


if __name__ == '__main__':
    main()