C++で二分木の対角和を求める方法

二分木の対角和とは

二分木の対角和とは、傾き -1 の平行な直線（対角線）で木を分割したとき、それぞれの対角線上に位置するノードのデータをすべて合計した値のことです。本記事では、C++ を使って再帰的な木の走査と std::map を組み合わせて、各対角線ごとの合計を効率的に求める方法を解説します。

ノード構造体の定義

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

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

ノードを生成する関数

次に、createNode(int data) 関数を作成します。この関数は int 型の値を受け取り、新しいノードを生成したうえでその値を data メンバに代入し、生成したノードへのポインタを返します。

Node * createNode(int data){
    Node * node = new Node;
    node->data = data;
    return node;
}

対角和を計算する再帰関数

中心となるのが diagonal_sum(Node *root, int depth, map<int, int> &diagonalSum) 関数です。引数としてルートノード、現在の深さ、そして参照渡しで受け取る diagonalSum マップを取ります。

処理の流れは以下のとおりです。

• ノードが NULL でなければ、現在のノードのデータを diagonalSum マップの「現在の深さ」をキーとする要素に加算します。
• 左の子へ移動するときは深さを +1 し、右の子へ移動するときは深さをそのまま維持して再帰呼び出しを行います。

これにより、同じ対角線上にあるノードは必ず同じ深さのキーに集約され、マップの各キーが対角線ごとの合計を保持することになります。

void diagonal_sum(Node *root, int depth, map<int, int> &diagonalSum){
    if(root){
        diagonalSum[depth]+=root->data;
        diagonal_sum(root->leftChild, depth+1, diagonalSum);
        diagonal_sum(root->rightChild, depth, diagonalSum);
    }
}

main関数での木の構築

main 関数内では、createNode(data) メソッドを使ってサンプル用の二分木を構築し、結果を格納するための map<int,int> 型の sumMap を用意します。ルートノード、初期深さ 1、そして sumMap を diagonal_sum に渡すことで、マップに「深さ（対角線番号）→ 合計値」のペアが格納されていきます。あとで走査できるよう、イテレータ it も宣言しておきます。

int main(){
    Node *root = createNode(1);
    root->rightChild = createNode(3);
    root->rightChild->leftChild = createNode(4);
    root->rightChild->leftChild->leftChild = createNode(12);
    root->rightChild->leftChild->rightChild = createNode(7);
    root->leftChild = createNode(2);
    root->leftChild->leftChild = createNode(9);
    root->leftChild->rightChild = createNode(6);
    root->leftChild->leftChild->rightChild = createNode(10);
    root->leftChild->rightChild->leftChild = createNode(11);
    root->rightChild->rightChild = createNode(5);
    map<int,int> sumMap;
    diagonal_sum(root, 1, sumMap);
    map<int,int>::iterator it;

結果の出力

最後に、for ループの中でイテレータ it を使って sumMap を先頭から末尾まで走査し、各対角線の合計値をタブ区切りで出力します。

for(it=sumMap.begin(); it!=sumMap.end();++it){
    int value = it->second;
    cout<<value<<"\t";
}

完全な実装例

ここまでの内容をまとめた、二分木の対角和を求めるプログラムの完全なコードが以下です。

#include<iostream>
#include<map>
using namespace std;
struct Node{
    int data;
    struct Node* leftChild, *rightChild;
};
Node * createNode(int data){
    Node * node = new Node;
    node->data = data;
    return node;
}
void diagonal_sum(Node *root, int depth, map<int, int> &diagonalSum){
    if(root){
        diagonalSum[depth]+=root->data;
        diagonal_sum(root->leftChild, depth+1, diagonalSum);
        diagonal_sum(root->rightChild, depth, diagonalSum);
    }
}
int main(){
    Node *root = createNode(1);
    root->rightChild = createNode(3);
    root->rightChild->leftChild = createNode(4);
    root->rightChild->leftChild->leftChild = createNode(12);
    root->rightChild->leftChild->rightChild = createNode(7);
    root->leftChild = createNode(2);
    root->leftChild->leftChild = createNode(9);
    root->leftChild->rightChild = createNode(6);
    root->leftChild->leftChild->rightChild = createNode(10);
    root->leftChild->rightChild->leftChild = createNode(11);
    root->rightChild->rightChild = createNode(5);
    map<int,int> sumMap;
    diagonal_sum(root, 1, sumMap);
    map<int,int>::iterator it;
    for(it=sumMap.begin(); it!=sumMap.end();++it){
        int value = it->second;
        cout<<value<<"\t";
    }
    return 0;
}

実行結果

上記のコードを実行すると、次のような出力が得られます（各数値はタブ区切りで表示されます）。

9   19   42

この結果は、深さ 1 の対角線（1 + 3 + 5 = 9）、深さ 2 の対角線（2 + 6 + 4 + 7 = 19）、深さ 3 の対角線（9 + 10 + 11 + 12 = 42）というように、各対角線上のノードの合計を表しています。