競プロ初心者向け。
先日参加したAtcoderにて、DPの計算量を削減する方法を知り、とても感動したのでメモ。
問題
一列に並んだ Nマスから成るマス目があり、マスには左から順番に1,2,…,Nの番号がついています。
このマス目で暮らしている高橋君は、現在マス 1にいて、後述の方法で移動を繰り返してマスNへ行こうとしています。
10以下の整数 Kと、共通部分を持たないK個の区間 [L1,R1],[L2,R2],…,[LK,RK]
が与えられ、これらの区間の和集合をSとします。
ただし、区間 [l,r]は l以上 r以下の整数の集合を表します。
- マスiにいるとき、Sから整数を 1つ選んで (dとする)、マス i+dに移動する。ただし、マス目の外に出るような移動を行ってはならない。
高橋君のために、マスNに行く方法の個数を 998244353で割った余りを求めてください。
ナイーブな解き方
一見すると、以下のような単純なDPで解けそう！という風に見えます。
dp[i]: 位置iに到達する移動方法の数
dp[i] = \sum_{s \in S}dp[i-s]
しかし、これだとdpテーブルの更新数(N) x Sの要素数(N) でN^2の計算量がかかってしまい、TLEエラーで通すことはできません。
高速化①
Sの要素が区間として与えられており、かつ区間同士で共通部分を持たないことに注目すると、
区間[Li, Ri]での移動はDP上のある区間の和と考えられます。
例: 区間[1, 5]の場合、 dp[i] = dp[i-1]+dp[i-2] + ... + dp[i-5]
これを利用すると、DPテーブルは下記のように書き換えられます。
dp[i] = \sum_{lr \in S}\sum_{j \in [l,r]}(dp[i-j])
(lr:Sに含まれる一つの区間)
高速化②
高速化①で、区間和を愚直に求めようとするとやはり計算量Nのため、全体で計算量N^2のままです。そこで、累積和の考え方を使うことで、区間和の計算を高速化します。
累積和とは、数列上で最初の要素から特定の要素までの値を順々に足し合わせたもので、
累積和を利用することで、数列上の区間和を下記のように計算できます。
(数列Lのi~jまでの区間和) = ([jまでの累積和] - [iまでの累積和])
この方法によって、任意の区間和を定数時間で計算することができ、全体の計算量をNKに抑えることができます。
累積和をsdpで表すと、DPは下記のようになります。
dp[i] = \sum_{lr \in S}(sdp[i-lr[r] - sdp[i-lr[l])
実装
pythonによる実装です。
dpの初期値(dp[0])は最初の場所へ行く方法の数なので1です。
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テーブル
sdp = [0] * (N+1) # dpテーブルの累積和
# 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)
まとめ
区間を利用したDPの計算と、累積和を利用した高速化という2つの概念が学べる良問でした。
D問題となると、単純なDPが実装できるだけではなく、計算量への意識が求められる、ということなのですね。