Pythonで解く「葉の値から構成する最小コスト二分木」問題 ― メモ化再帰による動的計画法

問題の概要

正の整数からなる配列 arr が与えられたとき、次の条件をすべて満たす二分木を考えます。

• 各ノードは、子を 0 個または 2 個持つ。
• 配列 arr の値は、木の中間順巡回（inorder traversal）における各葉の値に対応する。
• 各非葉ノードの値は、左部分木と右部分木それぞれにおける最大の葉の値の積と等しい。

考えられるすべての二分木の中から、各非葉ノードの値の合計が最小となるものを見つけるのが目的です。例えば、入力 arr = [6, 2, 4] の場合、出力は 32 になります。この配列からは次の 2 通りの木が構成できます。

上の図では、非葉ノードの値（24 と 8）の合計が 32 となり、これが達成可能な最小値です。

解法のアプローチ：区間DP＋メモ化再帰

この問題は、区間を分割しながら最小コストを求める動的計画法（メモ化再帰）で効率的に解けます。手順は以下の通りです。

1. 計算済みの結果をキャッシュするための辞書 memo を用意します。
2. 区間の両端 i と j を引数にとる関数 dp(i, j) を定義します。
3. j <= i の場合（区間に葉が1つしかない場合）は 0 を返します。
4. (i, j) がすでに memo に存在する場合は、memo[(i, j)] をそのまま返します。
5. 変数 res を無限大で初期化します。
6. k を i から j - 1 まで動かしながら、区間を左右に分割します。
   • res = min(res, dp(i, k) + dp(k+1, j) + max(arr[i:k+1]) * max(arr[k+1:j+1])) で更新します。
7. memo[(i, j)] = res として結果を保存し、それを返します。
8. 本体では dp(0, len(arr) - 1) を呼び出して答えを取得します。

ここでのポイントは、区間 [i, k] と [k+1, j] を結合するときのコストが、それぞれの区間内の最大葉の値の積になるという点です。この積が非葉ノードの値として加算されます。

実装例

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

class Solution(object):
   def mctFromLeafValues(self, arr):
      """
      :type arr: List[int]
      :rtype: int
      """
      self.memo = {}
      def dp(i, j):
         if j <= i:
            return 0
         if (i, j) in self.memo:
            return self.memo[(i, j)]
         res = float('inf')
         for k in range(i, j):
            res = min(res, dp(i, k) + dp(k + 1, j) + (max(arr[i:k + 1]) * max(arr[k + 1:j + 1])))
         self.memo[(i, j)] = res
         return self.memo[(i, j)]
      return dp(0, len(arr) - 1)

計算量

区間の組み合わせは O(n²)、各区間について分割位置 k の探索に O(n) かかるため、時間計算量は O(n³)、メモ用の空間計算量は O(n²) となります。

入力例

[6,2,4]

出力例

32