二分木の2つのノード間の距離を求めるクエリ – C++でのO(log n)手法

この問題では、二分木とQ個のクエリが与えられます。私たちのタスクは、C++でO(log n)の計算量を使って、二分木の2つのノード間の距離を求めるプログラムを作成することです。

問題の概要

各クエリでは、二分木の2つのノードが与えられ、その2つのノード間の距離を求める必要があります。ここでの「距離」とは、一方のノードからもう一方のノードに到達するために通過する必要がある辺（エッジ）の数を意味します。

具体例を見て問題を理解しましょう。

入力：二分木

クエリ数 = 3

Q1 -> [2, 6]

Q2 -> [4, 1]

Q3 -> [5, 3]

出力：3, 2, 3

解決アプローチ

この問題を解くには、最小共通祖先（LCA：Lowest Common Ancestor）と各ノードからの距離を利用した距離公式を使用します。

Distance(n1, n2) = distance(root, n1) + distance(root, n2) - 2 * distance(root, LCA)

この公式は、「ルートからn1までの距離」と「ルートからn2までの距離」を足し合わせると、ルートからLCAまでの経路が重複して2回カウントされるという性質を利用しています。問題を解くために、以下の手順に従います。

• 各ノード（n1、n2、LCA）のレベル（深さ）を求めます。

• オイラーツアー（Euler's Tour）に基づいて、二分木の配列を作成します。

• LCAを高速に求めるために、セグメント木（Segment Tree）を構築します。

アルゴリズムのポイント

オイラーツアーとは、木をDFS（深さ優先探索）で巡回し、各頂点に出入りするたびにその頂点を記録していく手法です。これにより、LCAは「2つのノードの最初の出現位置の間で、最も浅いレベルのノード」として区間最小値（RMQ）問題に帰着できます。セグメント木を使えば、この区間最小値クエリをO(log n)で処理できるため、大量のクエリにも効率的に対応できます。

実装例（C++）

#include <bits/stdc++.h>
#define MAX 1000
using namespace std;
int eulerArray[MAX];
int eIndex = 0;
int vis[MAX];
int L[MAX];
int H[MAX];
int level[MAX];
struct Node {
    int data;
    struct Node* left;
    struct Node* right;
};
struct Node* newNode(int data) {
    struct Node* temp = new struct Node;
    temp->data = data;
    temp->left = temp->right = NULL;
    return temp;
}
void FindNodeLevels(struct Node* root) {
    if (!root)
        return;
    queue<pair<struct Node*, int> > q;
    q.push({ root, 0 });
    pair<struct Node*, int> p;
    while (!q.empty()) {
        p = q.front();
        q.pop();
        level[p.first->data] = p.second;
        if (p.first->left)
            q.push({ p.first->left, p.second + 1 });
        if (p.first->right)
            q.push({ p.first->right, p.second + 1 });
    }
}
void createEulerTree(struct Node* root) {
    eulerArray[++eIndex] = root->data;
    if (root->left) {
        createEulerTree(root->left);
        eulerArray[++eIndex] = root->data;
    }
    if (root->right) {
        createEulerTree(root->right);
        eulerArray[++eIndex] = root->data;
    }
}
void creareEulerArray(int size) {
    for (int i = 1; i <= size; i++) {
        L[i] = level[eulerArray[i]];
        if (vis[eulerArray[i]] == 0) {
            H[eulerArray[i]] = i;
            vis[eulerArray[i]] = 1;
        }
    }
}
pair<int, int> seg[4 * MAX];
pair<int, int> min(pair<int, int> a, pair<int, int> b) {
    if (a.first <= b.first)
        return a;
    else
        return b;
}
pair<int, int> buildSegTree(int low, int high, int pos) {
    if (low == high) {
        seg[pos].first = L[low];
        seg[pos].second = low;
        return seg[pos];
    }
    int mid = low + (high - low) / 2;
    buildSegTree(low, mid, 2 * pos);
    buildSegTree(mid + 1, high, 2 * pos + 1);
    seg[pos] = min(seg[2 * pos], seg[2 * pos + 1]);
}
pair<int, int> LCA(int qlow, int qhigh, int low, int high, int pos) {
    if (qlow <= low && qhigh >= high)
        return seg[pos];
    if (qlow > high || qhigh < low)
        return { INT_MAX, 0 };
    int mid = low + (high - low) / 2;
    return min(LCA(qlow, qhigh, low, mid, 2 * pos), LCA(qlow, qhigh,mid + 1, high, 2 * pos +1));
}
int CalcNodeDistance(int node1, int node2, int size) {
    int prevn1 = node1, prevn2 = node2;
    node1 = H[node1];
    node2 = H[node2];
    if (node2 < node1)
        swap(node1, node2);
    int lca = LCA(node1, node2, 1, size, 1).second;
    lca = eulerArray[lca];
    return level[prevn1] + level[prevn2] - 2 * level[lca];
}
int main() {
    int N = 6;
    Node* root = newNode(1);
    root->left = newNode(2);
    root->right = newNode(3);
    root->left->left = newNode(4);
    root->left->right = newNode(5);
    root->right->left = newNode(6);
    FindNodeLevels(root);
    createEulerTree(root);
    creareEulerArray(2 * N - 1);
    buildSegTree(1, 2 * N - 1, 1);
    int Q = 4;
    int query[Q][2] = {{1, 5}, {4, 6}, {3, 4}, {2, 4} };
    for(int i = 0; i < Q; i++)
        cout<<"The distance between two nodes of binary tree is "<<CalcNodeDistance(query[i][0], query[i][1], 2 * N - 1)<<endl;
    return 0;
}

出力結果

The distance between two nodes of binary tree is 2
The distance between two nodes of binary tree is 4
The distance between two nodes of binary tree is 3
The distance between two nodes of binary tree is 1

処理の流れのまとめ

このコードでは、まずBFS（幅優先探索）によって各ノードのレベルを計算します。次にオイラーツアーでオイラー配列を構築し、各ノードの最初の出現位置（H配列）と対応するレベル（L配列）を記録します。そして、レベルの区間最小値クエリに答えるためのセグメント木を構築します。各クエリに対しては、2つのノードの最初の出現位置の間の区間からLCAを求め、距離公式に代入することでO(log n)で答えを計算できます。