atcoder ABC 179 D - Leaping Tak

問題

D - Leaping Tak

解法

素直にDPで解けそうだが、愚直にnに対して全ての移動を試すと $ O(N{^2}) $ となり $ 2≤N≤2×10{^5} $ の制約下でTLEしてしまう。

// このコードはTLEする
#include <bits/stdc++.h>
using namespace std;
#define COUT(x) cout<<(x)<<"\n"
#define REP(i, n) for(int i=0;i<n;i++)
#define endl "\n"
using ll = long long;
const int MOD = 998244353;

int main() {
  int n,k; cin >> n >> k;
  vector<int> l(k),r(k);
  REP(i,k) cin >> l[i] >> r[i];
  vector<ll> dp(n+1); // dp[マスi] = マスiに到達できる場合の数

  dp[1] = 1; // マス1から始まるため dp[1] = 1 にする
  for (int i = 2; i <= n; i++) {
    REP(j,k) {
      for(int k = l[j]; k <= r[j]; k++) {
        if (i-k < 0) continue;
        dp[i] += dp[i-k];
        dp[i] %= MOD;
      }
    }
  }
  COUT(dp[n]);
  return 0;
}

上記ではTLEするため、計算量を落とす方法を考える。

dpの累積和を取り、累積和の区間差分 dpsum[left] - dpsum[right] を取ることでその範囲の場合の数が得られることを利用し、計算量を落とす。

int main() {
  int n,k; cin >> n >> k;
  vector<int> l(k),r(k);
  REP(i,k) cin >> l[i] >> r[i];
  vector<ll> dp(n+1), dpsum(n+1);

  dp[1] = 1; // マス1から始まるため dp[1] = 1 にする
  dpsum[1] = 1; // dpの累積和

  for (int i = 2; i <= n; i++) {
    REP(j,k) {
      int li = max(i-r[j], 1);
      int ri = i-l[j];
      if (ri < 0) continue;

      dp[i] += dpsum[ri]-dpsum[li-1]; // 累積和の差からleft-right間の場合の数を得る
      dp[i] = dp[i] < 0 ? dp[i]+MOD : dp[i]%MOD; // MOD計算
    }
    dpsum[i] = (dpsum[i-1]+dp[i]) % MOD;
  }
  COUT(dp[n]);
  return 0;
}