C++で二分木の2つのノード間の距離を求める方法

問題の概要

いくつかのノードを持つ二分木が与えられているとします。このとき、2つのノード u と v の間の「距離」、つまり一方のノードからもう一方のノードへ移動する際に通る辺（エッジ）の本数を求めることを考えます。

例として、次のような二分木を扱います。

            1
           / \
          2   3
         / \ / \
        4  5 6  7
              \
               8

この木において、ノード (4, 6) 間の距離は 4（経路：4 → 2 → 1 → 3 → 6）、ノード (5, 8) 間の距離は 5（経路：5 → 2 → 1 → 3 → 6 → 8）となります。

解決のアプローチ：LCA（最小共通祖先）を利用する

2つのノード間の距離を効率的に求めるには、LCA（Lowest Common Ancestor：最小共通祖先）を利用するのが定石です。LCAとは、2つのノード両方の子孫となるノードの中で、最も根に近いノードのことです。

2つのノード間の距離は、次の式で表せます。

距離(u, v) = 距離(LCA, u) + 距離(LCA, v)

つまり、以下の手順で解くことができます。

1. 2つのノード u と v のLCAを求める。
2. LCAからノード u への距離を計算する。
3. LCAからノード v への距離を計算する。
4. 2つの距離を合計して答えとする。

C++での実装例

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

#include<iostream>
using namespace std;

class Node {
public:
    int data;
    Node *left, *right;
};

// 新しいノードを生成する補助関数
Node* getNode(int data) {
    Node* node = new Node;
    node->data = data;
    node->left = node->right = NULL;
    return node;
}

// 2つのノード n1, n2 の最小共通祖先（LCA）を再帰的に求める
Node* LowestCommonAncestor(Node* root, int n1, int n2) {
    if (root == NULL)
        return root;
    // 現在のノードがいずれかの対象ノードなら、それがLCA候補
    if (root->data == n1 || root->data == n2)
        return root;

    Node* left = LowestCommonAncestor(root->left, n1, n2);
    Node* right = LowestCommonAncestor(root->right, n1, n2);

    // 左右両方からNULL以外が返れば、現在のノードがLCA
    if (left != NULL && right != NULL)
        return root;
    if (left != NULL)
        return left;
    return right;
}

// ノード k の、root からの深さ（レベル）を求める
int getLevel(Node* root, int k, int level) {
    if (root == NULL) return -1;
    if (root->data == k) return level;

    int left = getLevel(root->left, k, level + 1);
    if (left == -1)
        return getLevel(root->right, k, level + 1);
    return left;
}

// 2つのノード a, b 間の距離を求める
int findDistance(Node* root, int a, int b) {
    Node* lca = LowestCommonAncestor(root, a, b);
    int dist1 = getLevel(lca, a, 0);
    int dist2 = getLevel(lca, b, 0);
    return dist1 + dist2;
}

int main() {
    Node* root = getNode(1);
    root->left = getNode(2);
    root->right = getNode(3);
    root->left->left = getNode(4);
    root->left->right = getNode(5);
    root->right->left = getNode(6);
    root->right->right = getNode(7);
    root->right->left->right = getNode(8);

    cout << "Distance between (4, 6) is: " << findDistance(root, 4, 6);
    cout << "\nDistance between (8, 5) is: " << findDistance(root, 8, 5);
    return 0;
}

実行結果

Distance between (4, 6) is: 4
Distance between (8, 5) is: 5

コードの解説

1. LowestCommonAncestor 関数

木を再帰的にたどりながら、対象となる2つの値 n1、n2 を探索します。左右の部分木のそれぞれから結果が返ってきた場合、そのノードが2つのノードを分岐させる共通の祖先であるため、そのノードをLCAとして返します。

2. getLevel 関数

指定したノード k が、基準ノード（ここではLCA）から何レベル下にあるかを再帰的に調べます。見つからない場合は -1 を返し、左部分木で見つからなければ右部分木を探索します。これにより、LCAから各ノードへの距離（辺の本数）が得られます。

3. findDistance 関数

LCAを求め、そこからノード a および b までの距離をそれぞれ計算し、その合計を返します。これが2つのノード間の距離になります。

計算量について

• 時間計算量： O(n)。LCAの探索とレベルの計算は、それぞれ最悪の場合で木の全ノードを訪問するためです。
• 空間計算量： O(h)。再帰呼び出しによるスタック領域が、木の高さ h に依存します（最悪の場合、木が偏っていると O(n) になります）。

まとめ

二分木における2つのノード間の距離は、「最小共通祖先（LCA）を求め、LCAからそれぞれのノードへの距離を足し合わせる」というシンプルな方針で効率よく計算できます。木構造の問題ではLCAが応用される場面が非常に多いため、この実装パターンはぜひ覚えておきましょう。