【ABC220】D - FG operation

atcoder.jp

解法

全探索しようとするとO(2^N)になってしまうような問題はDPで解けることが多い。
なので、DPで解くことを考えてみる。

i回目の操作F, Gは次のように言い換えられる。

• 左端と左からi + 1番目の値を計算して、現在の左端と入れ替える

ここで、i回目の操作を行った後に、i + 1回目の操作に影響するのは、i回目現在の左端の数字だけである。
(左からi + 1番目の値は最初から決まっているので)
そのため、現在の左端の数字を状態として持つDPをすることにする。

DPの定義としては、
・ dp[i][j] ← i回目の操作をした結果、左端の数字がjであるような場合の数

遷移は、操作F or 操作Gの二択になるので、

• 操作Fをするとき
  • (現在の左端の数字 + A[i + 1]) を10で割った数が次の左端の数字になるので、
  • dp[i + 1][(j + A[i + 1]) mod 10] += dp[i][j] mod 998244353
• 操作Gをするとき
  • (現在の左端の数字 × A[i + 1]) を10で割った数が次の左端の数字になるので、
  • dp[i + 1][(j × A[i + 1]) mod 10] += dp[i][j] mod 998244353

初期値は、左端の数字が最初のパターンになるので、
・dp[0][A[0]] = 1

答えは、N-1回操作した後に左端がjであるような数字をそれぞれ出力すればいいので、下のように実装すればよい。

for (int j = 0; j < 10; j++) {
  // dp[N-1][j]を出力
  cout << dp[N-1][j] << endl;
}

実装

modintを使ったので、mod 998244353の部分は省略しています。

void solve(long long N, std::vector<long long> A) {
    // DPテーブル
    vector<vector<mint>> dp(N, vector<mint>(10, 0));
    dp[0][A[0]] = 1;

    REP (i, N-1) {
        REP (j, 10) {
            // 操作Fをした場合
            dp[i + 1][(j + A[i + 1]) % 10] += dp[i][j];
            // 操作Gをした場合
            dp[i + 1][(j * A[i + 1]) % 10] += dp[i][j];
        }
    }

    // 出力
    REP (j, 10) {
        c(dp[N-1][j].val);
    }
}