Pythonで0からnの値で形成できる一意な二分探索木の個数を求めるプログラム

ある整数 n が与えられたとき、[0, n)(0 以上 n 未満)の範囲の数値を使って生成できる一意な二分探索木(BST)の個数を求めることを考えます。答えが非常に大きくなる可能性があるため、結果は 10^9 + 7 で割った余りを返します。

たとえば、入力が n = 3 の場合、出力は 5 になります。これは {0, 1, 2} の3つの値から作れる二分探索木の形状がちょうど5通り存在するためです。

この問題の鍵となる「カタラン数」

二分探索木の個数は、キーの具体的な値には依存せず、ノードの個数 n だけで決まります。n 個のノードから構成できる二分探索木の総数は、数学では「カタラン数」として知られており、次の式で表されます。

Catalan(n) = C(2n, n) ÷ (n + 1)

ここで C(2n, n) は二項係数(2n 個の中から n 個を選ぶ組み合わせの総数)です。つまりこの問題は、カタラン数を剰余演算のもとで高速に計算する問題に帰着します。

解法のアプローチ

剰余を取りながらの除算は直接行えないため、フェルマーの小定理を利用したモジュラ逆数を使います。手順は以下のとおりです(m = 10^9 + 7 とします)。

• 分子 numer を 1、分母 denom を n + 1 で初期化する
• i を 1 から n までループする
  • numer ← numer × (n + i) を計算し、m で割った余りを取る
  • denom ← denom × i を計算し、m で割った余りを取る
• ループ終了後、denom のモジュラ逆数 denom^(m−2) mod m を numer に掛ける
• numer mod m を結果として返す

この方法なら、ループ部分は O(n)、最後のべき乗計算は O(log m) で済むため、n が大きくなっても高速に答えを求められます。

補足：なぜ denom^(m−2) で割り算ができるのか

m が素数のとき、フェルマーの小定理により a^(m−1) ≡ 1 (mod m) が成り立ちます。これを変形すると a^(m−2) ≡ a^(−1) (mod m)、つまり a の m−2 乗は「a で割る操作」と同じ働きをします。10^9 + 7 は素数なので、この性質によって除算を乗算に置き換えられます。

実装例(Python)

class Solution:
    def solve(self, n):
        m = 10 ** 9 + 7
        numer = 1
        denom = n + 1
        for i in range(1, n + 1):
            numer *= n + i
            numer %= m
            denom *= i
            denom %= m
        numer *= pow(denom, m - 2, m)
        return numer % m

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

入力

4

出力

14

n = 4 の場合、C(8, 4) ÷ 5 = 70 ÷ 5 = 14 となり、0 から 3 までの4つの値で作れる一意な二分探索木が 14 通りであることが確認できます。

計算量の目安

• 時間計算量:O(n + log m)
• 空間計算量:O(1)