🕌

ABC220 D - FG operation C++解答例

2021/09/28に公開

C++

AtCoder Beginner Contest 220 D - FG operationをC++で解きます。

問題

問題文をAtCoderのページより引用します。

問題文

0以上9以下の整数からなる長さ $N$ の数列 $A = (A_1, \ldots, A_N)$ があり、この順に左から右に並んでいます。
数列の長さが1になるまで、操作 $F$ または操作 $G$ を繰り返し行います。

操作 $F$ ：左端の2つの値( $x, y$ とする)を削除した後、一番左に $(x + y) \% 10$ を挿入する

操作 $G$ ：左端の2つの値( $x, y$ とする)を削除した後、一番左に $(x \times y) \% 10$ を挿入する

なお、 $a \% b$ は $a$ を $b$ で割った余りを意味します。
$K = 0, 1, \dots, 9$ について、以下の問題に答えてください。

操作手順としてあり得るものは $2^{N - 1}$ 通りありますが、このうち最終的に残る値が $K$ となる操作手順は何通りありますか？
ただし答えは非常に大きくなる可能性があるので、998244353で割った余りを求めてください。

制約

$2 \leq N \leq 10^5$

$0 \leq A_i \leq 9$

入力はすべて整数

解答例

動的計画法を用いる

問題文にある通り、操作手順の全ての場合の数は $2^{N - 1}$ 通り存在するため、全てを全探索することは到底困難です。
別の解き方を考えてみましょう。

ここで、与えられた2つの操作の内容を観察すると、どちらの操作も「数列の左端の2つの要素に対して演算をした後、その2つの要素と演算結果を置換する」操作であることがわかります。
また、1回の操作を行うために必要な情報は、現在の数列の左端にある2つの要素のみであることもわかります。

すなわち、それまでに行った操作の結果に関係なく、今ある数列の左端の要素だけが現在の操作の結果に影響するということです。
このような問題は動的計画法を用いて解くことができそうです。

DPテーブルを設計する

DPテーブルの構造

動的計画法で解けそうなことがわかったので、次にDPテーブルを設計します。
DPテーブルを以下のように定義します。

$dp[i][j]$ ： $i$ 回目の操作を行った後に、数列の左端の要素が $j$ になったときの、操作手順の場合の数

ただし、 $0 \leq i \leq N - 1$ 、 $0 \leq j \leq 9$ です^[1]。

DPの遷移

次に遷移を考えます。今回の動的計画法は、 $i - 1$ 回目までの操作の結果が $i$ 回目の操作の演算結果に影響を及ぼすので、遷移の形は「貰う遷移」に似たものになります。

$i - 1$ 回目の操作を終えた段階で、数列の左端の要素の値が $j$ になった時を考えます。
そのような操作手順の場合の数は $dp[i - 1][j]$ に格納されています。
このとき、 $i$ 回目の操作で「操作 $F$ を行う」のと「操作 $G$ を行う」の2通りが考えられます。
2つの操作のうちのどちらかを行う際、演算の対象となるのは数列の左端の要素です。
今、一番左端にある要素の値は $j$ です。そして、その1つ右の要素は、数列 $A$ を0-indexedで考えたとき、 $A_{i}$ になります^[2]。

以上を踏まえると、 $i$ 回目の操作の演算結果は、操作 $F$ 、操作 $G$ のそれぞれについて、以下のようになります。

操作 $F$ ： $(j + A_{i}) \% 10$
操作 $G$ ： $(j \times A_{i}) \% 10$

$i - 1$ 回目の操作を行って数列の左端の要素が $j$ になったような場合の数が $dp[i - 1][j]$ なので、その状態で $i$ 回目の操作を行ったときの場合の数は、次のように更新されます。

操作 $F$ ： $dp[i][(j + A_{i}) \% 10] = dp[i][(j + A_{i}) \% 10] + dp[i - 1][j]$
操作 $G$ ： $dp[i][(j \times A_{i}) \% 10] = dp[i][(j \times A_{i}) \% 10] + dp[i - 1][j]$

尚、DPテーブルを更新した際に998244353の余りを取るのを忘れないように注意してください。

DPテーブルの初期条件

今回の動的計画法ではDPテーブルの値として場合の数をとっているので、DPテーブルの各要素の初期値は0にしておきます。
また、まだ操作を行っていない状態(「0回目の操作を行った後の状態」と表現します)のとき、数列の左端の要素の値は $A_0$ です。
「何も操作を行わない」ことも1通りと数えられるので、DPテーブルの $dp[0][A_0]$ には1を設定しておきます。

プログラムの実装例

以上を踏まえてプログラムを実装します。設計した通りにDPテーブルを実装すればよいでしょう。

d.cpp

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
 
int main() {
  int64_t N;
  std::cin >> N;
  std::vector<int64_t> A(N);
  for (int64_t i = 0; i < N; i++) {
    std::cin >> A.at(i);
  }

  // DPテーブル。各要素の初期値は0にする。
  std::vector<std::vector<int64_t>> dp(N, std::vector<int64_t>(10, 0));
  // 初期条件
  dp.at(0).at(A.at(0)) = 1;
  // DPループ
  for (int64_t i = 1; i < N; i++) {
    for (int64_t j = 0; j < 10; j++) {
      if (dp.at(i - 1).at(j) > 0) {
        dp.at(i).at((j + A.at(i)) % 10) += dp.at(i - 1).at(j);
        dp.at(i).at((j * A.at(i)) % 10) += dp.at(i - 1).at(j);
 
        dp.at(i).at((j + A.at(i)) % 10) = dp.at(i).at((j + A.at(i)) % 10) % 998244353;
        dp.at(i).at((j * A.at(i)) % 10) = dp.at(i).at((j * A.at(i)) % 10) % 998244353;
      }
    }
  }
 
  for (int64_t i = 0; i < 10; i++) {
    std::cout << dp.at(N - 1).at(i) << std::endl;
  }
 
  return 0;
}

実際に提出したコードはこちら。

脚注

要素数が $N$ の数列に対して操作 $F$ と操作 $G$ は合計で $N - 1$ 回まで行われるためです。また、初期条件としてまだ操作をしていない(「0回目の操作」と表現します)ときの数列の左端の要素が必要となるので、 $i = 0$ のときの情報もDPテーブルに含めるようにしています。 ↩︎
数列 $A$ を0-indexedで考えると、 $A = (A_0, A_1, \ldots, A_{N - 1})$ となります。1回目の操作の演算対象は $A_0$ と $A_1$ で、2回めの操作の演算対象は、1回目の操作の演算結果および $A_2$ です。従って、 $i$ 回目の操作の演算対象は、 $i - 1$ 回目の操作の演算結果と、 $A_{i}$ ということになります。 $i$ は $N - 1$ までと定義されているので、配列の領域外参照を起こす心配はありません。 ↩︎

問題

問題文

制約

解答例