Official

D - シフトの割り当て / Shift Assignment Editorial by physics0523


この問題は、動的計画法 (DP) により解くことができます。

以下の値を持っておくことを考えます。

  • \(dp[k] = \{\) 現在までのシフトの決め方であって、直近で \(k\) 日連続で A のシフトに入っているようなものの数 \(\}\)

このうち、 \(dp[0]\) および \(dp[1]\) だけ覚えておけばこの問題においては十分です。

この \(dp\) を、 \(S\) を前から \(1\) 文字ずつ見ていきながら更新します。
\(ndp\) をひとつ先の \(dp\) であるとします。

  • 現在見ている文字が A または C であるなら、それを A のシフトにすることができる。
    • \(ndp[1]\) に \(dp[0]\) を加算する。
    • \(dp[1]\) の状態から A のシフトを入れると問題文中の条件に違反します。
  • 現在見ている文字が B または C であるなら、それを B のシフトにすることができる。
    • \(ndp[0]\) に \(dp[0]\) および \(dp[1]\) を加算する。
      • B のシフトに入るので、 A のシフトの連続はリセットされます。

最終的な \(dp[0]+dp[1]\) が答えとなります。

本解法の時間計算量は \(O(N)\) です。

実装例 (C++):

#include<bits/stdc++.h>

using namespace std;
const int mod=998244353;

int main(){
  int N;
  string S;
  cin >> N >> S;
  array<int,2> dp={1,0};
  for(auto &nx : S){
    array<int,2> ndp={0,0};
    if(nx!='B'){
      ndp[1]+=dp[0]; ndp[1]%=mod;
    }
    if(nx!='A'){
      ndp[0]+=dp[0]; ndp[0]%=mod;
      ndp[0]+=dp[1]; ndp[0]%=mod;
    }
    dp=ndp;
  }
  cout << (dp[0]+dp[1])%mod << "\n";
  return 0;
}

posted:
last update: