C++で母音の順列を数える：行列累乗法による効率的な解法

問題概要

整数 n が与えられたとき、次のルールに従って構成できる長さ n の文字列が何通りあるかを数えます。

• 使用できるのは小文字の母音（a, e, i, o, u）のみ
• 'a' の後ろには 'e' のみ続けられる
• 'e' の後ろには 'a' または 'i' のみ続けられる
• 'i' の後ろに 'i' は続けられない
• 'o' の後ろには 'i' または 'u' のみ続けられる
• 'u' の後ろには 'a' のみ続けられる

答えは非常に大きな値になる可能性があるため、10⁹ + 7 で割った余りを求めます。

入力例と出力例

たとえば入力が 2 の場合、出力は 10 となります。条件を満たす文字列は次の 10 通りだからです。

"ae", "ea", "ei", "ia", "ie", "io", "iu", "oi", "ou", "ua"

解法の考え方：隣接行列と行列累乗

この問題は、各母音をノード、遷移ルールをエッジとみなしたグラフの経路数え上げとして捉えることができます。まず 5×5 の隣接行列 A を定義し、「A[i][j] = 1 ならば母音 i の後ろに母音 j を続けられる」という意味を持たせます。

長さ n の文字列の総数は、Aⁿ⁻¹ の全要素の合計に等しくなります。n が大きい場合でも、行列累乗（繰り返し二乗法）を使えば O(log n) 回程度の行列乗算で高速に計算できます。また、加算・乗算のたびに 10⁹ + 7 で剰余を取ることで、オーバーフローを防ぎながら正確な値を維持できます。

アルゴリズムの手順

1. 法 m = 10⁹ + 7 を定義する
2. 関数 add(a, b)：((a mod m) + (b mod m)) mod m を返す
3. 関数 mul(a, b)：((a mod m) × (b mod m)) mod m を返す
4. 関数 solve(n)：
   
   • 5×5 の隣接行列 A を {{0,1,0,0,0},{1,0,1,0,0},{1,1,0,1,1},{0,0,1,0,1},{1,0,0,0,0}} として定義する
   • 5×5 の単位行列 result を用意する
   • n を 1 減らし、その回数だけ result ← result × A を繰り返し計算する
   • result の全要素の合計を答えとして返す

C++ 実装例

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
const lli m = 1e9+7;
lli add(lli a, lli b){
   return ((a%m) + (b%m))%m;
}
lli mul(lli a, lli b){
   return ((a%m) * (b%m))%m;
}
class Solution {
   public:
   void multiply(lli A[5][5], lli B[5][5]){
      lli C[5][5];
      for(lli i =0;i<5;i++){
         for(lli j=0;j<5;j++){
            lli temp =0;
            for(lli k =0;k<5;k++){
               temp = add(temp,mul(A[i][k],B[k][j]));
            }
            C[i][j] = temp;
         }
      }
      for(lli i =0;i<5;i++){
         for(lli j =0;j<5;j++){
            A[i][j] = C[i][j];
         }
      }
   }
   lli solve(lli n){
      lli A[5][5] = { { 0, 1, 0, 0, 0 }, { 1, 0, 1, 0, 0 }, { 1, 1,
      0, 1, 1 }, { 0, 0, 1, 0, 1 }, { 1, 0, 0, 0, 0 } };
      lli result[5][5];
      for (lli i = 0; i < 5; i++) {
         for (lli j = 0; j < 5; j++) {
            if (i == j)
               result[i][j] = 1;
            else
               result[i][j] = 0;
         }
      }
      n--;
      for (int i = 1; i <= n; i++)
      multiply(result, A);
      lli sum = 0;
      for (lli i = 0; i < 5; i++) {
         for (lli j = 0; j < 5; j++) {
            sum = add(result[i][j], sum);
         }
      }
      return sum;
   }
   int countVowelPermutation(int n) {
      return solve(n);
   }
};
main(){
   Solution ob;
   cout << (ob.countVowelPermutation(2));
}

実行結果

入力：

2

出力：

10

まとめ

本記事では、母音の遷移ルールに従う長さ n の文字列の総数を数える問題を取り上げました。遷移ルールを隣接行列として表現し、行列累乗によって Aⁿ⁻¹ を求めることで、単純な全探索では処理しきれない大きな n に対しても効率的に答えを得られます。剰余演算を適切に挟むことで、巨大な数値によるオーバーフローも回避できる点がポイントです。