Problem p02662 (python)

Fixed submission:
import sys

sys.setrecursionlimit(10 ** 7)
input = sys.stdin.readline
f_inf = float('inf')
mod = 998244353


def resolve():
    """ 愚直解 """
    dp = [[0] * (s + 1) for _ in range(n + 1)]
    dp[0][0] = 1
    for i in range(1, n + 1):
        a = A[i - 1]
        for prev in range(i):
            for j in range(s + 1):
                dp[i][j] += dp[prev][j]
                if a + j > s:
                    continue
                dp[i][a + j] += dp[prev][j]
                dp[i][a + j] %= mod
    print(dp)
    res = 0
    for i in range(1, n + 1):
        res += dp[i][-1]
        res %= mod
    print(res)


def resolve2():
    """ DP + 累積和 """
    dp = [[0] * (s + 1) for _ in range(n + 1)]
    dp[0][0] = 1
    cnt_num = [0] * (s + 1)
    cnt_num[0] = 1
    res = 0
    for i in range(1, n + 1):
        a = A[i - 1]
        for j in range(s + 1):
            dp[i][j] += cnt_num[j]
            dp[i][j] %= mod
            if a + j > s:
                continue
            dp[i][a + j] += cnt_num[j]
            dp[i][a + j] %= mod
        for k in range(s + 1):
            cnt_num[k] = cnt_num[k] + dp[i][k]
            cnt_num[k] %= mod
        res += dp[i][-1]
        res %= mod
    print(res)


if __name__ == '__main__':
    n, s = map(int, input().split())
    A = list(map(int, input().split()))
    resolve2()

Buggy submission (verdict: Runtime Error):
import sys

sys.setrecursionlimit(10 ** 7)
input = sys.stdin.readline
f_inf = float('inf')
mod = 998244353


class SegTree:
    """
    init(init_val, ide_ele): 配列init_valで初期化 O(N)
    update(k, x): k番目の値をxに更新 O(N)
    query(l, r): 区間[l, r)をsegfuncしたものを返す O(logN)
    """

    def __init__(self, init_val, segfunc, ide_ele):
        """
        init_val: 配列の初期値
        segfunc: 区間にしたい操作
        ide_ele: 単位元
        n: 要素数
        num: n以上の最小の2のべき乗
        tree: セグメント木(1-index)
        """
        n = len(init_val)
        self.segfunc = segfunc
        self.ide_ele = ide_ele
        self.num = 1 << (n - 1).bit_length()
        self.tree = [ide_ele] * 2 * self.num
        # 配列の値を葉にセット
        for i in range(n):
            self.tree[self.num + i] = init_val[i]
        # 構築していく
        for i in range(self.num - 1, 0, -1):
            self.tree[i] = self.segfunc(self.tree[2 * i], self.tree[2 * i + 1])

    def update(self, k, x):
        """
        k番目の値をxに更新
        k: index(0-index)
        x: update value
        """
        k += self.num
        self.tree[k] = x
        while k > 1:
            self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
            k >>= 1

    def query(self, left, right):
        """
        [left, right)のsegfuncしたものを得る
        left: index(0-index)
        right: index(0-index)
        """
        res = self.ide_ele
        left += self.num
        right += self.num
        while left < right:
            if left & 1:
                res = self.segfunc(res, self.tree[left])
                left += 1
            if right & 1:
                res = self.segfunc(res, self.tree[right - 1])
            left >>= 1
            right >>= 1
        return res


def resolve():
    """ 愚直解 """
    dp = [[0] * (s + 1) for _ in range(n + 1)]
    dp[0][0] = 1
    for i in range(1, n + 1):
        a = A[i - 1]
        for prev in range(i):
            for j in range(s + 1):
                dp[i][j] += dp[prev][j]
                if a + j > s:
                    continue
                dp[i][a + j] += dp[prev][j]
                dp[i][a + j] %= mod
    print(dp)
    res = 0
    for i in range(1, n + 1):
        res += dp[i][-1]
        res %= mod
    print(res)


def resolve2():
    """ DP + セグ木 """
    dp = [[0] * (s + 1) for _ in range(n + 1)]
    dp[0][0] = 1
    cnt_num = [SegTree([0] * n, lambda x, y: x + y, 0) for _ in range(s + 1)]
    cnt_num[0].update(0, 1)
    for i in range(1, n + 1):
        a = A[i - 1]
        for j in range(s + 1):
            dp[i][j] += cnt_num[j].query(0, i)
            if a + j > s:
                continue
            dp[i][a + j] += cnt_num[j].query(0, i)
            dp[i][a + j] %= mod
        for k in range(s + 1):
            cnt_num[k].update(i, dp[i][k])

    res = 0
    for i in range(1, n + 1):
        res += dp[i][-1]
        res %= mod
    print(res)


if __name__ == '__main__':
    n, s = map(int, input().split())
    A = list(map(int, input().split()))
    resolve2()