第11回日本情報オリンピック予選(過去問) D - パスタ(Pasta) Python解答例

第11回日本情報オリンピック予選(過去問) D - パスタ(Pasta)をPythonで解きます。

問題文

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

問題

あなたはパスタが大好物であり，毎日，晩御飯にパスタを作って食べている．あなたはトマトソース，クリームソース，バジルソースの$3$種類のパスタを作ることができる．

$N$日間の晩御飯の予定を考えることにした．それぞれの日に$3$種類のパスタから$1$種類を選ぶ．ただし，同じパスタが続くと飽きてしまうので，$3$日以上連続して同じパスタを選んではいけない．また，$N$日のうちの$K$日分のパスタはすでに決めてある．

入力として$N$の値と，$K$日分のパスタの情報が与えられたとき，条件をみたす予定が何通りあるかを$10\,000$で割った余りを求めるプログラムを作成せよ．

入力

入力は$K + 1$行からなる．

$1$行目には$2$つの整数$N, K$($3 \leqq N \leqq 100$，$1 \leqq K \leqq N$) が空白を区切りとして書かれている．

$1 + i$行目 ($1 \leqq i \leqq K$) には$2$つの整数$A_i, B_i$($1 \leqq A_i \leqq N$，$1 \leqq B_i \leqq 3$) が空白を区切りとして書かれている．これは，$A_i$日目のパスタはすでに決まっており，$B_i = 1$のときはトマトソースであり，$B_i = 2$のときはクリームソースであり，$B_i = 3$のときはバジルソースであることを表す．$A_i$($1 \leqq i \leqq K$) は全て異なる．与えられる入力データにおいて，条件をみたす予定は$1$通り以上あることが保証されている．

解答例

動的計画法で解きます。

DPテーブルの定義

以下のようにDPテーブルを定義します。ただし、$1 \leq i \leq N, 0 \leq j \leq 2$です。

$dp[i][j][0]$：$i - 1$日目に$j$を食べずに$i$日目に$j$を食べる場合の数
$dp[i][j][1]$：$i - 1$日目にも$j$を食べて$i$日目に$j$を食べる場合の数

初期条件は以下のようになります。ここで、すでにパスタの種類が決まっている日の集合を$\mathbf{A}$、日付$a \in \mathbf{A}$に対応するパスタの種類を表す数値を返す写像を$f(a)$とおきます。

$$\begin{cases} \begin{cases} dp[1][j][0] &= 1 \quad (1 \notin \mathbf{A}, 0 \leq j \leq 2) \\ dp[1][j][1] &= 0 \quad (1 \notin \mathbf{A}, 0 \leq j \leq 2) \end{cases} \\ \begin{cases} dp[1][f(1)][0] &= 1 \quad (1 \in \mathbf{A}) \\ dp[1][f(1)][1] &= 0 \quad (1 \in \mathbf{A}) \end{cases} \\ \end{cases}$$

遷移式は以下のようになります。尚、$j$以外の2つの種類をそれぞれ$j'$、$j''$と表記することとします。

$$\begin{cases} dp[i][j][0] &= \sum^2_{k=1} dp[i - 1][j'][k] + \sum^2_{k=1} dp[i - 1][j''][k] \\ dp[i][j][1] &= dp[i - 1][j][0] \end{cases} \quad (1 \leq i \leq N, 0 \leq j \leq 2)$$

考察パート

まず、「3日以上連続して同じパスタを選んではいけない」ルールが無い状態を考えてみます。予めパスタの種類も決まっていないとします。
このときの場合の数の考え方は極めて単純に考えることができます。以下のようなグラフを考え、各ノードに至るパスの本数を数えればよいです。
この図では横軸に$i$、縦軸に$j (1 \leq j \leq 3)$を取り、各ノード内の数字は$i$日目に種類$j$のパスタを食べる場合の数を表しています。$i$日目に種類$j$のパスタを食べることに当たるノードを$(i, j)$と表現することとします。

次に、3日以上連続して同じパスタを食べてはいけないルールを課してみます。すると、このままのグラフではルールを考慮することができないことに気づきます。
たとえば、3日目に種類1のパスタを食べる場合の数を数えてみます。上の場合と同じように、2日目に種類1, 2, 3のパスタを食べた場合の数の和を求めれば良いように見えますが、ノード$(2, 1)$は「1日目に種類1のパスタを食べる場合の数」の情報も含んでしまっています。
従ってこれを計上してしまうとルールに反することになってしまいます。一方で、ノード$(2, 1)$は「1日目に種類1以外のパスタを食べる場合の数」の情報を含んでいます。これは数えないといけません。

以上のことから、上の図のようなグラフの形式では問題の制約をうまく考慮できないことがわかりました。
これを解決するために、各ノードに以下の状態を持たせることを考えます。

• 前日に同じ種類のパスタを食べたかどうか。

これを考慮したグラフは以下の図のようになります。

まず、$i - 1$日目に食べた種類とは異なる種類のパスタ$j$を$i$日目に食べることを考えます。この場合、$j$は$i$日目で連続しないので、$i - 1$日目に$j'$、$j''$を食べる場合の数の総数を取れば良いです。

次に、$i - 1$日目に$j$を食べ、更に$i$日目にも$j$を食べる場合を考えます。このとき、$i$日目に$j$を食べられるのは、$i - 2$日目に$j$以外を食べていた場合のみです。(そうでなければ3日連続で$j$を食べることになってしまいます。)
さて、$i - 2$日目に$j$以外を食べ、$i - 1$日目に$j$を食べる場合の数は、ノード$(i - 1, j)$の$[0]$に保存されています。従って、これを$(i, j)$の[1]に持ってくれば良いです。

この計算を各日付と各種類に対して実施することで、全ての場合の数を数え上げることができます。

最後に、すでに食べる種類が決まっている日付について考えます。$A_i$日目に種類$B_i$のパスタを食べることが決まっているということは、それ以外の種類のパスタを食べる場合の数は0であるということです。
これを表現するには、単純にノード$(A_i, B_{i}'), (A_i, B_{i}'')$の値を0にしてしまえば良いです。
上の計算の後、$i$日目の予定がすでに決まっていたことがわかったら、$dp[i][j]$以外の$dp[i]$のリストを0で初期化することでこの操作を表現することができます。

実装例

実際のコード例を以下に示します。
集合$\mathbf{A}$とその写像は、Pythonの辞書で表現します。これにより、$i$日目の予定がすでに決まっているかどうか($i$が集合$\mathbf{A}$に含まれているかどうか)とその写像$f(a)$の取得の操作を$\mathcal{O}(1)$で行うことができます。

d.py

from typing import List, Dict
import itertools


def main():
    # 入力受け取り
    N, K = map(int, input().split())
    A, B = map(list, zip(*[list(map(int, input().split())) for i in range(K)]))

    # 法
    mod: int = 10000

    # すでに決まっている予定の情報を格納しておくための辞書
    D: Dict[int, int] = {a: b - 1 for a, b in zip(A, B)}

    # DPテーブル
    dp: List[List[List[int]]] = [[[0] * 2 for j in range(3)] for i in range(N + 1)]

    # 初日の予定が決まっているかどうかで初期条件が変わるので場合分けして初期化
    if 1 in D:
        dp[1][D[1]] = [1, 0]
        dp[1][(D[1] + 1) % 3] = [0, 0]
        dp[1][(D[1] + 2) % 3] = [0, 0]
    else:
        dp[1] = [[1, 0] for i in range(3)]

    # DP遷移
    for i in range(2, N + 1):
        for j in range(3):
            # 前日と違う種類を食べる場合
            dp[i][j][0] = sum(dp[i - 1][(j + 1) % 3]) + sum(dp[i - 1][(j + 2) % 3])
            # 前日と同じ種類を食べる場合
            dp[i][j][1] = dp[i - 1][j][0]

            # 余りを取る
            dp[i][j][0] %= mod
            dp[i][j][1] %= mod

        # i日目の予定がすでに決まっていた場合
        if i in D:
            # 決まっている種類以外のパスタを食べる場合の数を0にする
            # TIPS: 0, 1, 2の数字のうちある数字以外の数字は、1、2をそれぞれ
            # 足した余りを取ることで求められる。
            dp[i][(D[i] + 1) % 3] = [0, 0]
            dp[i][(D[i] + 2) % 3] = [0, 0]

    # 結果の出力
    # dp[N]はリストのリストになっていて、最終結果は各リストの要素の和である。
    # itertools.chainは与えられたシーケンスの中身を全て展開して一つのシーケンス
    # に纏める働きを持つので、dp[N]の中身を展開して渡している。
    print(sum(itertools.chain(*dp[N])) % mod)


if __name__ == "__main__":
    main()

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