Pythonで二分木の「良い」葉ノードペアの数を求めるプログラム

問題の概要

二分木と整数値 d が与えられます。異なる2つの葉ノードからなるペアのうち、両ノード間の最短経路の長さが d 以下であるものを「良いペア(good pair)」と呼びます。この記事では、Pythonを使って木の中に良いペアがいくつ存在するかを求める方法を解説します。

たとえば、次のような二分木を考えてみましょう。

この木に対して d = 4 とした場合、答えは 2 になります。(8, 7) と (5, 6) の2つのペアは経路長がどちらも 2 で d 以下だからです。一方、(7, 5) や (8, 6) などのペアは経路長が 5 になり、d = 4 を超えるため良いペアとして数えられません。

解法のアプローチ：後順走査による再帰

この問題は、木を後順(ポストオーダー)で再帰的に走査するアプローチで効率よく解けます。各ノードでは「その部分木に含まれる葉ノードまでの距離」のリストを返し、左右の部分木をつなぐ位置でペアの成立を判定します。

具体的な手順は以下のとおりです。

• 答えを格納する変数 sol を 0 で初期化します。
• util(root) という関数を定義します。
• root が null(空)の場合は、空のリストを返します。
• root が葉ノードの場合は、[[0, 0]] を返します(2番目の要素が距離を表し、自分自身との距離は 0 です)。
• それ以外の場合は、次の処理を行います。
  • 左部分木の結果 l = util(root.left) と、右部分木の結果 r = util(root.right) を取得します。
  • l・r の各要素について、距離(2番目の要素)を 1 ずつ加算します。これは現在のノードを1つ上る分の距離を反映させるためです。
  • r の各ノード n と l の各ノード n1 のすべての組み合わせを調べ、n[1] + n1[1] ≤ d であれば sol を 1 増やします。
  • l と r を結合したリストを親ノードへ返します。
• メイン処理では util(root) を呼び出した後、sol を返します。

同じ部分木内にある葉同士のペアは、それらの共通祖先にあたる上位ノードで必ず判定されるため、この方法でもペアの重複や漏れなく数えられる点がポイントです。

Pythonによる実装例

それでは、実際のコードを見てみましょう。

class TreeNode:
   def __init__(self, val=0, left=None, right=None):
      self.val = val
      self.left = left
      self.right = right

class Solution:
   def __init__(self):
      self.sol = 0

   def solve(self, root, d):
      def util(root):
         if not root:
            return []
         if not root.left and not root.right:
            return [[0, 0]]
         else:
            l = util(root.left)
            r = util(root.right)
            for n in l:
               n[1] += 1
            for n in r:
               n[1] += 1
            for n in r:
               for n1 in l:
                  if n[1] + n1[1] <= d:
                     self.sol += 1
            return l+r
      util(root)
      return self.sol

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.right = TreeNode(4)
root.left.right.left = TreeNode(8)
root.left.right.right = TreeNode(7)
root.right.left = TreeNode(5)
root.right.right = TreeNode(6)

d = 4
ob = Solution()
print(ob.solve(root, d))

入力

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.right = TreeNode(4)
root.left.right.left = TreeNode(8)
root.left.right.right = TreeNode(7)
root.right.left = TreeNode(5)
root.right.right = TreeNode(6)
d = 4

出力

2

コードのポイント

• util() は、部分木内の各葉ノードへの距離を格納したリストを返します。各要素は [ダミー値, 距離] の形式で、実際に使われるのは2番目の距離の値です。
• 左右の部分木をまたぐペアだけをそのノードでカウントすることで、すべてのペアがちょうど1回ずつ(最小共通祖先の位置で)判定されます。
• 距離の加算(+1)は、再帰的に親へ戻る際に経路長を正しく更新するために行っています。

計算量

各ノードで左右のリストの全組み合わせを比較するため、最悪ケースの時間計算量は O(n²)(n はノード数)となります。空間計算量は、再帰スタックと距離リストの保持により O(n) です。