Pythonで二分木の複数ノードから最小共通祖先（LCA）を求めるプログラム

はじめに

二分木が与えられたとき、その中に含まれる複数のノードすべての最小共通祖先（Lowest Common Ancestor：LCA）を求めたい場面はよくあります。二分木における最小共通祖先とは、指定されたノード x1, x2, x3, …, xn のすべてを子孫にもつノードの中で、最も深い位置にあるノードのことです。なお、あるノード自身も自分自身の子孫とみなせる点に注意してください。

本記事では、木のルートノードと、祖先を求めたいノードのリストを入力として受け取り、該当するノードを返すプログラムをPythonで実装します。

問題の例

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

ここで、祖先を求めたいノードのリストが [6, 8] だった場合、出力は 7 になります。

これは、ノード 6 とノード 8 の両方を子孫にもつノードの中で最も深いものが 7 であるためです。

解法のアプローチ

この問題は、以下の手順で解くことができます。

1. 補助関数 fn() を定義します。引数としてノードを受け取ります。
   • ノードが null（None）の場合は、そのままノードを返します。
   • ノードが探索対象のノード集合に含まれている場合も、そのノードを返します。
   • 左部分木に対して fn(左の子) を、右部分木に対して fn(右の子) を再帰的に呼び出します。
   • 左右の結果がどちらも null でなければ、現在のノードが共通祖先であるため、そのノードを返します。
   • どちらか一方だけが null でない場合は、null でない方を返します。
2. 対象ノードのリストを作成し、各要素に対して search_node() を使って木の中から実際のノードオブジェクトを取得します。
3. 取得したノードを集合（set）に変換することで、以降の存在判定を高速に行えるようにします。
4. ルートノードに対して fn() を呼び出し、その結果を返します。

実装例

以下がPythonによる実装です。

import collections

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 is 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 print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.data, end=', ')
        print_tree(root.right)

def solve(root, node_list):
    # 対象ノードを実際のノードオブジェクトとして取得し、集合に変換
    targets = {search_node(root, elem) for elem in node_list}

    def fn(node):
        if not node:
            return node
        if node in targets:
            return node
        left = fn(node.left)
        right = fn(node.right)
        if left and right:
            return node
        return left or right

    return fn(root)

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

入力

make_tree([5, 3, 7, 2, 4, 6, 8]), [6, 8]

出力

7

処理の流れと計算量

まず solve() 内で、指定された値（6 と 8）に対応するノードオブジェクトを search_node() によって木から検索し、集合 targets に格納します。その後、fn() をルートから再帰的に実行すると、以下のように動作します。

• 葉に向かって深さ優先で探索を進め、対象ノード（6 や 8）に到達した時点でそのノードを上位に返します。
• 左右両方の部分木から非nullの結果が返ってきた最初のノードが、まさに最小共通祖先となります。この例ではノード 7 が該当します。

計算量は、各ノードを一度ずつ訪問するため O(n)（n はノード数）、追加のメモリは再帰の深さ分の O(h)（h は木の高さ）となります。

まとめ

二分木の最小共通祖先を求める問題は、再帰的な後順走査（post-order traversal）を活用することで簡潔に解くことができます。対象ノードを事前に集合化しておくことで判定も効率化でき、実務でも応用範囲の広いテクニックです。ぜひ自分のコードでも試してみてください。