Pythonで二分木の最小共通祖先（LCA）を求める方法

二分木が与えられたとき、指定した2つのノードの最小共通祖先（Lowest Common Ancestor：LCA）を求める問題を考えてみましょう。ノード p と q の LCA とは、p と q の両方を子孫として持つノードの中で、最も深い位置にあるノードのことです。

例えば、二分木が [3,5,1,6,2,0,8,null,null,7,4] という形式で表されている場合、木の構造は次のようになります。

この場合、ノード 5 と ノード 1 の LCA は 3 となります。

解法のアプローチ

この問題は、再帰を使って次の手順で解くことができます。

• 木が空（None）の場合は、None を返します。
• 現在のノード（root）の値が p または q と一致する場合は、そのノードを返します。
• 左部分木に対して再帰的に LCA を求め、結果を left に格納します。
• 右部分木に対して再帰的に LCA を求め、結果を right に格納します。
• left と right がどちらも None でない場合、p と q が左右別々の部分木に存在することを意味するため、現在のノードが LCA となり、root を返します。
• それ以外の場合は、None でない方（left または right）を返します。

それでは、実際の実装を見て理解を深めましょう。

実装例

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
class Solution(object):
    def lowestCommonAncestor(self, root, p, q):
        if not root:
            return None
        if root.data == p or root.data ==q:
            return root
        left = self.lowestCommonAncestor(root.left, p, q)
        right = self.lowestCommonAncestor(root.right, p, q)
        if right and left:
            return root
        return right or left
ob1 = Solution()
tree = make_tree([3,5,1,6,2,0,8,None,None,7,4])
print(ob1.lowestCommonAncestor(tree, 5, 1).data)

入力

[3,5,1,6,2,0,8,null,null,7,4]
5
1

出力

3

このアルゴリズムの計算量は、木のすべてのノードを最大1回ずつ訪問するため O(n) です（n はノード数）。また、再帰の深さは木の高さに依存するため、バランスの取れた木であれば O(log n)、最悪ケース（偏った木）では O(n) のスタック領域が必要になります。