Pythonで母音の遷移規則に従って作成できる文字列の数をカウントするプログラム

数 n が与えられたとき、以下の規則に従って生成できる長さ n の文字列の総数を求めることを考えます。

• 各文字は小文字の母音 [a, e, i, o, u] のいずれかである
• 「a」の後に続けられるのは「e」のみ
• 「e」の後に続けられるのは「a」または「i」
• 「i」の後に「i」を続けることはできない
• 「o」の後に続けられるのは「i」または「u」
• 「u」の後に続けられるのは「a」のみ

結果が非常に大きくなる可能性があるため、答えは 10^9 + 7 で割った余りを返します。

例として、入力が n = 2 の場合、出力は 10 になります。このとき生成できる2文字の文字列は、["ae", "ea", "ei", "ia", "ie", "io", "iu", "oi", "ou", "ua"] の10通りです。

解き方

この問題は動的計画法(DP)を用いることで効率的に解けます。ポイントは、「各母音で終わる文字列が何通りあるか」を状態として管理し、遷移規則に基づいて次の長さの状態へ順に更新していくことです。

具体的な手順は以下の通りです。

• モジュール値 m = 10^9 + 7 を設定する
• n が 0 の場合は 0 を返す
• 5つの変数 a, e, i, o, u を用意し、すべて 1 で初期化する(長さ1の文字列は各母音について1通りずつ存在するため)
• n - 1 回のループを実行し、各回で次の遷移を行う
  • a := e + i + u(「e」「i」「u」の後に「a」を付けられるため)
  • e := a + i(「a」「i」の後に「e」を付けられるため)
  • i := e + o(「e」「o」の後に「i」を付けられるため)
  • o := i(「i」の後に「o」を付けられるため)
  • u := i + o(「i」「o」の後に「u」を付けられるため)
• 最後に (a + e + i + o + u) mod m を返す

なお、Pythonではタプルの同時代入を利用すると、右辺がすべて評価されてから一括で左辺に代入されるため、一時変数を用意せずに状態更新を簡潔に記述できます。

それでは、以下の実装例を見て理解を深めましょう。

実装例

class Solution:
    def solve(self, n):
        m = (10 ** 9 + 7)
        if n == 0:
            return 0
        a = e = i = o = u = 1
        for _ in range(n-1):
            a, e, i, o, u = e+i+u, a+i, e+o, i, i+o
        return (a + e + i + o + u) % m

ob = Solution()
print(ob.solve(3))

入力

3

出力

19

n = 3 の場合、19通りの文字列が生成できることが確認できます。この解法の時間計算量は O(n)、空間計算量は O(1) であり、n が非常に大きい場合でも高速に動作するのが特徴です。