Pythonでn個の異なるノードから生成できるBST（二分探索木）の数を求めるプログラム

整数 n が与えられたとします。[1, 2, ..., n] のような n 個の異なる値があるとき、これらの値を使って構成できるBST（二分探索木）の総数を数える必要があります。答えが非常に大きくなる可能性があるため、結果は 10^9+7 で割った余りとして返します。

たとえば、入力が n = 3 の場合、出力は 14 になります。

解法のアプローチ

この問題は、動的計画法（DP）を使って効率的に解くことができます。ある値を根に選ぶと、それより小さい値で作られる左部分木と、大きい値で作られる右部分木に分割できるため、小さな部分問題の答えを組み合わせることで全体の答えが求まります。

具体的には、次の手順に従います。

• a := [0, 1] で初期化されたリストを用意する
• m := 10^9+7（答えの剰余を取るための定数）
• max_n := 1000（事前計算しておく最大のノード数）
• k を 2 から max_n + 1 まで繰り返し、次の処理を行う
  • a の末尾に (1 + Σ a[i] × a[k−i]（i は 1 以上 k 未満のすべての値）) mod m を追加する
• (a[n + 1] − 1) mod m を答えとして返す

実装例

理解を深めるために、以下のPythonコードを見てみましょう。

def solve(n):
   a = [0, 1]
   m = 10**9+7
   max_n = 1000

   for k in range(2, max_n + 2):
      a.append((1 + sum(a[i] * a[k - i] for i in range(1, k))) % m)
   return ((a[n + 1] - 1) % m)

n = 3
print(solve(n))

入力

3

出力

14

このように、漸化式を用いた動的計画法によって、n 個の異なるノードから構成できるBSTの数を効率よく求めることができます。あらかじめ max_n まで計算しておくことで、複数のクエリにも高速に対応できる点もポイントです。