PythonでAjob言語の単語から部分列を選択する方法の数を求めるプログラム

問題の概要

ここでは、「Ajob言語」という奇妙な言語を考えます。この言語には無限個の文字が存在します。私たちはこの言語のn個の単語を知っており、1番目の単語は1文字、2番目の単語は2文字、3番目の単語は3文字……というように、i番目の単語の長さはちょうどi文字になっています。さらに、各単語を構成する文字はすべて互いに異なります。

このn個の単語の中から任意の1つを選び、その部分列(元の並びの一部を抜き出した列)を作ることを考えます。ただし、部分列の長さは元の単語の長さよりkだけ短くなければなりません。つまり、選んだ単語の長さをLとすると、部分列の長さは(L − k)です。長さがk未満の単語は選ぶことができません。

2つの部分列が「異なる」のは、長さが異なる場合か、同じ位置に異なる文字が含まれる場合です。求める答えは、選べる部分列の総数を素数pで割った余りです。

例えば、入力が n = 6、k = 5、p = 11 のとき、出力は 7 になります。

解法のアプローチ

長さLの単語から長さ(L − k)の部分列を作る方法の数は、二項係数 C(L, k) に等しくなります。したがって、選べる単語(長さがk以上のもの)全体について C(L, k) を足し合わせればよいことになります。

ここで役立つのが「ホッケースティック恒等式」です。これにより、合計 Σ C(L, k)(L = k から n まで)= C(n+1, k+1) と簡潔に表せます。コード内で n と k をそれぞれ1増やしているのはこのためです。

あとは C(n+1, k+1) mod p を計算するだけですが、n や k が非常に大きな値になる可能性もあるため、ルーカスの定理を利用します。これは、二項係数の剰余をp進法の各桁ごとの二項係数の積として計算できるという定理です。また、剰余演算での割り算のために、モジュラ逆数と逆階乗のテーブルを事前に構築し、memo辞書でキャッシュして再利用します。

アルゴリズムの手順

• 結果をキャッシュするための空の辞書 memo を用意する
• n := n + 1、k := k + 1 とする(ホッケースティック恒等式のため)
• fact := 要素 1 のみを持つリスト(階乗テーブル)
• i を 1 から p−1 まで動かしながら、fact の末尾に(fact の最後の要素 × i mod p)を追加する
• memo に p が存在すれば、inv_fact := memo[p] を使用する
• 存在しない場合は、以下を実行する
  • inv := 要素 0 と 1 を持つリスト(モジュラ逆数テーブル)
  • i を 2 から p−1 まで動かしながら、inv の末尾に(p − (p÷i の商) × inv[p mod i] mod p)を追加する
  • inv_fact := 要素 1 のみを持つリスト
  • i を 1 から p−1 まで動かしながら、inv_fact の末尾に(inv_fact の最後の要素 × inv[i] mod p)を追加する
  • memo[p] := inv_fact としてキャッシュする
• ret := 1 と初期化する
• n > 0 の間、以下を繰り返す
  • n1 := n mod p、k1 := k mod p
  • k1 > n1 であれば 0 を返す(その桁では二項係数が 0 になるため)
  • ret := ret × fact[n1] × inv_fact[k1] × inv_fact[n1 − k1] mod p
  • n := n÷p の商、k := k÷p の商
• ret を返す

実装例(Python)

理解を深めるために、以下の実装を見てみましょう。

memo = {}
def solve(n, k, p):
   n += 1
   k += 1
   fact = [1]
   for i in range(1, p):
      fact.append(fact[-1] * i % p)
   if p in memo:
      inv_fact = memo[p]
   else:
      inv = [0, 1]
      for i in range(2, p):
         inv.append(p - p // i * inv[p % i] % p)
      inv_fact = [1]
      for i in range(1, p):
         inv_fact.append(inv_fact[-1] * inv[i] % p)
      memo[p] = inv_fact
   ret = 1
   while n > 0:
      n1 = n % p
      k1 = k % p
      if k1 > n1:
         return 0
      ret = ret * fact[n1] * inv_fact[k1] * inv_fact[n1 - k1] % p
      n //= p
      k //= p
   return ret

n = 6
k = 5
p = 11
print(solve(n, k, p))

入力

6, 5, 11

出力

7

解説のポイント

このプログラムの計算量は、前処理(階乗・逆階乗テーブルの構築)に O(p)、ルーカスの定理の適用に約 O(log_p n) であり、n や k が巨大な整数でも高速に動作します。モジュラ逆数は「p − p÷i × inv[p mod i] mod p」という漸化式で効率的に求められ、逆階乗も同様に累積的に構築できます。memo 辞書によるキャッシュのおかげで、同じ素数 p に対する問い合わせが複数回発生しても、逆階乗テーブルの再計算を回避でき、処理をさらに高速化できます。