Pythonで二分木の2つのノード間の距離を求めるプログラム

二分木が与えられたとき、その中の2つのノード間の距離を求めることを考えます。グラフの場合と同じように、2つのノードを結ぶ経路上の辺(エッジ)の数を数え、その本数を距離として返します。

二分木のノード構造

木の各ノードは、次のような構造を持っています。

data : <整数値>
right : <木の別のノードへのポインタ>
left : <木の別のノードへのポインタ>

問題の例

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

この木において、ノード「2」とノード「8」の間の距離を求めたいとします。このときの出力は 4 になります。

ノード2からノード8へ至る経路上の辺は、(2, 3)、(3, 5)、(5, 7)、(7, 8) の4本です。経路に含まれる辺が4本あるため、2つのノード間の距離は4となります。

解き方のアプローチ

この問題を解くためには、次の手順に従います。

1. 関数 findLca() を定義する(引数: root、p、q)
   • root が None であれば、None を返します。
   • root.data が p または q と一致していれば、root を返します。
   • left := findLca(root.left, p, q)
   • right := findLca(root.right, p, q)
   • left と right がどちらも None でなければ、root を返します。
   • それ以外の場合は、left または right を返します。
2. 関数 findDist() を定義する(引数: root、data)
   • queue := 新しい deque(両端キュー)を作成します。
   • (root, 0) というペアをキューの末尾に追加します。
   • キューが空でない限り、次の処理を繰り返します。
     • current := キューの先頭ペアの1番目の値(ノード)
     • dist := キューの先頭ペアの2番目の値(距離)
     • current.data が data と一致すれば、dist を返します。
     • current.left が存在すれば、(current.left, dist+1) をキューに追加します。
     • current.right が存在すれば、(current.right, dist+1) をキューに追加します。
3. node := findLca(root, p, q)
4. findDist(node, p) + findDist(node, q) を返します。

ここでのポイントはLCA(最小共通祖先)です。2つのノード p と q の共通の祖先の中で最も深い位置にあるノードが LCA であり、p と q を結ぶ経路は必ず LCA を通ります。したがって、「LCA から p までの距離」と「LCA から q までの距離」をそれぞれ求めて足し合わせれば、p と q の間の距離が得られます。

実装例

理解を深めるために、以下の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 findLca(root, p, q):
    # 2つのノード p と q の最小共通祖先(LCA)を再帰的に求める
    if root is None:
        return None
    if root.data in (p, q):
        return root
    left = findLca(root.left, p, q)
    right = findLca(root.right, p, q)
    if left and right:
        return root
    return left or right

def findDist(root, data):
    # BFS(幅優先探索)で root から data までの距離を求める
    queue = collections.deque()
    queue.append((root, 0))
    while queue:
        current, dist = queue.popleft()
        if current.data == data:
            return dist
        if current.left:
            queue.append((current.left, dist + 1))
        if current.right:
            queue.append((current.right, dist + 1))

def solve(root, p, q):
    # LCAを基点に、p・qそれぞれへの距離の合計を返す
    node = findLca(root, p, q)
    return findDist(node, p) + findDist(node, q)

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

入力

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

出力

4

計算量について

findLca() は木を一度だけ再帰的に走査するため O(n)、findDist() も幅優先探索により最悪で O(n) の時間がかかります。したがって、全体の時間計算量は O(n) です。空間計算量についても、再帰スタックとキューの分だけ O(n) となります(n はノード数)。

このように、LCAを基点として左右それぞれの距離を求めて合計する手法を使えば、二分木上の任意の2ノード間の距離を効率的に計算できます。