[PYTHON] [Competition Pro] Speeding up DP using cumulative sum

For beginners in competition professionals. Atcoder, which I participated in the other day, learned how to reduce the amount of DP calculation, and was very impressed, so I made a note.

problem

Atcoder ABC 179 D-Leaping Tak

There are squares consisting of N squares arranged in a row, and the squares are numbered 1, 2,…, N in order from the left. Mr. Takahashi, who lives in this square, is currently in square 1 and is trying to go to square N by repeating the movement by the method described later. An integer K of 10 or less and K intervals [L1, R1], [L2, R2],…, [LK, RK] that have no intersection Is given, and let S be the union of these intervals. However, the interval [l, r] represents a set of integers greater than or equal to l and less than or equal to r.

--When you are in cell i, select an integer from S (let's say d) and move to cell i + d. However, you must not move out of the square.

For Takahashi, find the remainder of dividing the number of ways to go to Mass N by 998244353.

Naive solution

At first glance, it seems to be solved with the following simple DP! It looks like that.

dp [i]: number of ways to move to reach position i

dp[i] = \sum_{s \in S}dp[i-s]

However, in this case, the number of updates (N) x the number of elements (N) of the dp table requires the amount of calculation of N ^ 2, and it cannot be passed with a TLE error.

Speed up ①

Noting that the element of S is given as an interval and that the intervals have no intersection, Movement in the interval [Li, Ri] is considered to be the sum of certain intervals on the DP. Example: For the interval [1, 5], dp [i] = dp [i-1] + dp [i-2] + ... + dp [i-5] By using this, the DP table can be rewritten as follows.

dp[i] = \sum_{lr \in S}\sum_{j \in [l,r]}(dp[i-j])

(One section included in lr: S)

Speed up ②

If you try to find the sum of intervals in the speedup ①, the amount of calculation is N, so the total amount of calculation is N ^ 2. Therefore, by using the concept of cumulative sum, the calculation of interval sum is speeded up.

Cumulative sum is the sum of the values from the first element to a specific element on a sequence. By using the cumulative sum, you can calculate the interval sum on the sequence as follows.

(Sum of intervals from i to j of sequence L) = ([Cumulative sum up to j]-[Cumulative sum up to i])

With this method, any interval sum can be calculated in a constant time, and the total amount of calculation can be suppressed to NK. Expressing the cumulative sum in sdp, DP is as follows.

dp[i] = \sum_{lr \in S}(sdp[i-lr[r] - sdp[i-lr[l])

Implementation

It is an implementation by python. The initial value of dp (dp [0]) is 1 because it is the number of ways to get to the first place.

N, K = map(int, input().split())
LR = [list(map(int, input().split())) for _ in range(K)]
mod = 998244353

S = []
for k in range(K):
    for i in range(LR[k][0], LR[k][1] + 1):
        S.append(i)

dp = [0] * N #dp table
sdp = [0] * (N+1) #Cumulative sum of dp tables

#Set the initial value of DP
dp[0] = 1
sdp[1] = 1

for n in range(1, N):
    for lr in LR:
        left = max(0, n - lr[1])
        right = max(0, n - lr[0] + 1)
        dp[n] += sdp[right] - sdp[left]
        dp[n] %= mod
    sdp[n+1] = (sdp[n] + dp[n]) % mod

res = dp[N-1]
print(res)

Summary

It was a good question to learn two concepts: DP calculation using intervals and speeding up using cumulative sum. When it comes to the D problem, not only can you implement a simple DP, but you also need to be aware of the amount of calculation.

Recommended Posts

[Competition Pro] Speeding up DP using cumulative sum
[Python] Speeding up processing using cache tools
Notes on using dict in python [Competition Pro]
Speeding up numerical calculation using NumPy / SciPy: Application 1
Cumulative sum commentary