B - Nbonacci Sequence

2 secs 1024 MB

解説

題意通り $f(x)$ を計算すると、 $O(NX)$ の計算量が必要となり、間に合いません。
そこで、 $F(x) = f(0) + f(1) + \cdots + f(x) \ ($ 便宜上 $f(0) = 0$ とする $)$ の 累積和 を用いて、 $f(x) = F(x) - F(x-N-1)$ と計算し高速化することで、 $O(N+X)$ で解くことができます。

想定解

Python

xxxxxxxxxx
 
MOD = 998244353
​
N, X = map(int, input().split())
A = list(map(int, input().split()))
​
if X <= N:
    print(A[X-1])
    exit()
​
f = [0]*(X+1)
acc_f = [0]*(X+1)
​
for i, a in enumerate(A, start=1):
    f[i] = a
    acc_f[i] = (a + acc_f[i-1]) % MOD
​
for i in range(N+1, X+1):
    val = (acc_f[i-1] - acc_f[i-N-1]) % MOD
    f[i] = val
    acc_f[i] = (val + acc_f[i-1]) % MOD
​
print(f[X])