Pythonで二分木の最小共通祖先（LCA）を求めるプログラム

二分木と、その中の2つの特定ノード x と y が与えられたとします。このとき、2つのノードの最小共通祖先（Lowest Common Ancestor：LCA）を二分木から見つけ出す必要があります。

二分木における最小共通祖先とは、ノード x とノード y の両方の子孫となるノードの中で、最も深い位置にあるノードのことです。なお、あるノード自身も自分自身の子孫として扱える点に注意してください。求めたノードを出力として返します。

具体例

例えば、次のような二分木が与えられたとします。

ここで x = 2、y = 4 とした場合、出力は 3 になります。

ノード 2 とノード 4 の両方が子孫となっているノードは 3 であり、それより下位の共通の祖先は存在しないためです。

解き方の手順

この問題は、深さ優先探索（DFS）を用いて以下のように解くことができます。

• 関数 dfs() を定義します。引数としてノードを受け取ります。
  • ノードが null（None）の場合は何も返しません。
  • ノードがリスト [x, y] に含まれる場合：
    • 左部分木に対して dfs() を呼び出し、結果を left に格納します。
    • 右部分木に対して dfs() を呼び出し、結果を right に格納します。
    • left または right のどちらかが非ゼロ（見つかっている）であれば、現在のノードが答えとなります。ans にノードを代入し、そのノードを返します。
  • 上記以外の場合も、左右の子ノードに対してそれぞれ dfs() を再帰的に呼び出します。
  • left と right の両方が null でなければ、現在のノードが x と y の分岐点（＝最小共通祖先）です。ans にノードを代入し、そのノードを返します。
  • それ以外の場合は、left または right のうち null でない方を返します。
• ルートノードに対して dfs(root) を実行し、結果を ans に格納します。
• ans を返します。

実装コード

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

def insert(temp, data):
   que = []
   que.append(temp)
   while len(que):
      temp = que[0]
      que.pop(0)
      if not temp.left:
         if data is not None:
            temp.left = TreeNode(data)
         else:
            temp.left = TreeNode(0)
         break
      else:
         que.append(temp.left)
      if not temp.right:
         if data is not None:
            temp.right = TreeNode(data)
         else:
            temp.right = TreeNode(0)
         break
      else:
         que.append(temp.right)

def make_tree(elements):
   Tree = TreeNode(elements[0])
   for element in elements[1:]:
      insert(Tree, element)
   return Tree

def search_node(root, element):
   if root == None:
      return None
   if root.data == element:
      return root
   res1 = search_node(root.left, element)
   if res1:
      return res1
   res2 = search_node(root.right, element)
   return res2

def solve(root, x, y):
   def dfs(node):
      if not node:
         return
      if node in [x, y]:
         left = dfs(node.left)
         right = dfs(node.right)
         if left or right:
            ans = node
            return node
      left = dfs(node.left)
      right = dfs(node.right)
      if left and right:
         ans = node
         return node
      return left or right
   ans = dfs(root)
   return ans

root = make_tree([5, 3, 7, 2, 4, 1, 7, 6, 8, 10])
print(solve(root, search_node(root, 2), search_node(root, 4)).data)

入力

make_tree([5, 3, 7, 2, 4, 1, 7, 6, 8, 10]),
search_node(root, 2),
search_node(root, 4)

出力

3

計算量について

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量は O(n)、再帰呼び出しによる空間計算量は木の高さに依存し、最悪の場合 O(n) となります。二分木のサイズが大きくなっても効率的に動作する、シンプルかつ実用的な手法です。