けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 104 D - We Love ABC (400 点)

これ好きなん!!!
400 点問題が問題分読んで 5 秒以内に方針立ったのは珍しいので、成長かもなん。
そして、DEGwer さんや chokudai さんが「DP は割とコーディングしながら細部詰めてる」と言っていたので、それを実践してみて、確かに良さそうかもなん。
でも int 型で事足りるやろと考えてオーバーフローやらかしていたことになかなか気付かなかったのは、反省なのんな...そこで 5 分縮められるはずなん。。。

問題へのリンク

問題概要

"BCABBACCBCCACA"
のような A, B, C のみからなる文字列 S の「ABC 数」とは、

  • S[i] = 'A'
  • S[j] = 'B'
  • S[k] = 'C'
  • i < j < k

を満たすような (I, j, k) の組の個数のことである。今、

????C?????B??????A???????

のような A, B, C, ? からなる文字列が与えられる。? には A, B, C のどれを当てはめてもよい。考えられるすべての文字列についての「ABC 数」の総和を 1000000007 で割ったあまりで求めよ。

制約

  • 3 <= |S| <= 105

考えたこと

まずは「?」がなかったらどうするかを考えてみるのは多分よさそう。ナイーブには  O(|S|^{3}) かかるので、それを高速化するために、一瞬「累積和」とか「セグメントツリー」とかが脳裏をよぎりそうになる。でもそういうのは「?」が出て来るとちょっとつらいイメージがある。なんとなくだけど、こういう雰囲気の問題は速攻で DP なイメージがある・・・

それに、「?」がなかった場合は、

  • 文字列 S と "ABC" との共通部分列

に関する問題になる。こう考えると DP がとても自然に思えて来る。「?」があっても DP なら自然に対応することができる。

  • dp[ i + 1 ][ 0 ] := i 文字目まで考えたときに ABC との照合をスタートすらしていないようなものの個数
  • dp[ i + 1 ][ 1 ] := i 文字目まで考えたときに ABC との照合において A までは照合させたようなものの個数
  • dp[ i + 1 ][ 2 ] := i 文字目まで考えたときに ABC との照合において B までは照合させたようなものの個数
  • dp[ i + 1 ][ 3 ] := i 文字目まで考えたときに ABC との照合において C までは照合させたようなものの個数

という感じで DP を定義した。

  • 初期値: dp[0][0] = 1
  • 求める値: dp[S.size()][3]

遷移は以下のようにに場合分けして考えた:

i 文字目でABC カウントを進めない場合

  • i 文字目が '?' だったら 3 文字ありうるので、
    dp[ i + 1 ][ j ] += dp[ i ][ j ] * 3

  • i 文字目が '?' 以外だったら
    dp[ i + 1 ][ j ] += dp[ i ][ j ]

i 文字目で ABC カウントを進める場合

  • i 文字目が 'A' または '?' だった場合は「なし」から「A」にカウンタを勧められるので
    dp[ i + 1 ][ 1 ] += dp[ i ][ 0 ]

  • i 文字目が 'B' または '?' だった場合は「A」から「B」にカウンタを勧められるので
    dp[ i + 1 ][ 2 ] += dp[ i ][ 1 ]

  • i 文字目が 'C' または '?' だった場合は「B」から「C」にカウンタを勧められるので
    dp[ i + 1 ][ 3 ] += dp[ i ][ 2 ]

こんな感じで遷移を組み立てることができた。実際にはオーバーフローに注意しつつ実装すると AC をもらうことができた。

#include <iostream>
#include <string>
#include <cstring>
using namespace std;

const int MOD = 1000000007;
string S;

void add(long long &a, long long b) {
  a += b;
  if (a >= MOD) a -= MOD;
}

long long dp[210000][5];

int main() {
  cin >> S;
  memset(dp, 0, sizeof(dp));
  dp[0][0] = 1;
  for (int i = 0; i < S.size(); ++i) {
    for (int j = 0; j < 5; ++j) {
      if (S[i] != '?') add(dp[i+1][j], dp[i][j]);
      else add(dp[i+1][j], dp[i][j] * 3 % MOD);
    }
    if (S[i] == 'A' || S[i] == '?') add(dp[i+1][1], dp[i][0]);
    if (S[i] == 'B' || S[i] == '?') add(dp[i+1][2], dp[i][1]);
    if (S[i] == 'C' || S[i] == '?') add(dp[i+1][3], dp[i][2]);  
  }
  cout << dp[S.size()][3] << endl;
}

別解

A, B, C が 3 個しかないことを活かして、B の位置を決めて左右の A, C を数えるような方法もとれる