C++でO(n)の計算量で二分木の直径を求める方法【新手法】

はじめに

二分木（バイナリツリー）の直径とは、木の中で最も長い経路の長さを指します。各ノードに対して「左部分木の高さ + 右部分木の高さ + 1」を計算し、その最大値が直径となります。

本記事では、この性質を利用して、各ノードごとに「left_height + right_height + 1」を計算しながら結果を更新していく手法を紹介します。この方法なら、時間計算量はO(n)に抑えられます。

ツリーノード構造体の定義

まず、データ本体と左右の子ノードへのポインタを持つ、ツリーノードを表す構造体を定義します。最初に作成されたノードはルートノードとなり、それ以降に作成されるノードは子ノードとして扱われます。

struct Node {
    int data;
    struct Node *leftChild, *rightChild;
};

newNode関数の実装

次に、newNode(int data)関数を作成します。この関数はint型の値を受け取り、ノードのdataメンバに代入します。戻り値として、作成したNode構造体へのポインタを返します。また、新しく作成されたノードの左右の子はNULLに初期化されます。

struct Node* newNode(int data){
    struct Node* newNode = new Node;
    newNode->data = data;
    newNode->leftChild = newNode->rightChild = NULL;
    return (newNode);
}

diameter関数の実装

diameter(Node* root)関数は、ルートノードを受け取り、そのノードがNULLかどうかを判定します。その後、INT_MINで初期化された変数ansを定義します。height(root, ans)の戻り値は変数height_of_treeに格納され、最終的にansが関数から返されます。

int diameter(Node* root){
     if (root == NULL)
         return 0;
     int ans = INT_MIN;
     int height_of_tree = height(root, ans);
     return ans;
}

height関数の実装

height(Node* root, int& ans)関数は、ルートノードと参照渡しの変数ansを受け取ります。木に対して再帰的な走査を行い、各部分木の高さを計算します。再帰呼び出しのたびに、ansの最大値が第2引数として更新されていきます。ここで、ansは「ans と 1 + 左部分木の高さ + 右部分木の高さ」のうち大きい方の値になります。

サンプルコード

以下に、O(n)の手法で二分木の直径を求める完全な実装例を示します。

#include <iostream>
using namespace std;
struct Node {
    int data;
    Node* leftChild, *rightChild;
};
struct Node* newNode(int data){
    struct Node* newNode = new Node;
    newNode->data = data;
    newNode->leftChild = newNode->rightChild = NULL;
    return (newNode);
}
int height(Node* root, int& ans){
    if (root == NULL)
        return 0;
    int left_height = height(root->left, ans);
    int right_height = height(root->right, ans);
    ans = max(ans, 1 + left_height + right_height);
    return 1 + max(left_height, right_height);
}
int diameter(Node* root){
    if (root == NULL)
        return 0;
    int ans = INT_MIN;
    int height_of_tree = height(root, ans);
    return ans;
}
int main(){
    struct Node* root = newNode(1);
    root->left = newNode(2);
    root->right = newNode(3);
    root->left->left = newNode(4);
    root->left->right = newNode(5);
    printf("Diameter is %d\n", diameter(root));
    return 0;
}

実行結果

上記のコードを実行すると、次の出力が得られます。

Diameter is 4

まとめ

この手法では、高さの計算と直径の更新を1回の再帰走査で同時に行うため、全体の時間計算量はO(n)、空間計算量は再帰スタックの深さに依存してO(h)（hは木の高さ）となります。従来の各ノードで高さを再計算するO(n²)のアプローチと比べて大幅に効率的であり、大きな二分木でも高速に直径を求められるのが特徴です。