C++で二分木の最大垂直和を求める方法

はじめに

二分木が与えられたとき、垂直順序走査における各垂直列のノード値の合計を計算し、その中から最大値を求めて出力するのが本記事の課題です。例として、以下のような二分木を考えてみましょう。

この二分木を垂直順序走査すると、各列の合計は次のようになります。

4
2
1 + 5 + 6 = 12
3 + 8 = 11
7
9

各列の合計の中で最大となるのは 12 です。

アルゴリズムの考え方

アプローチはシンプルです。幅優先探索(BFS)を用いて垂直順序走査を行い、各ノードに水平距離を割り当てます。ルートの水平距離を 0 とし、左に移動するごとに -1、右に移動するごとに +1 とします。同じ水平距離を持つノードは同じ垂直列に属するため、map を使って水平距離ごとにノードの値をグループ化します。走査が完了したら、各グループの合計を計算し、その最大値を返します。

手順のまとめ

1. ルートを水平距離 0 でキューに追加する。
2. キューが空になるまで、先頭のノードを取り出し、そのキーを現在の map に記録する。
3. 左の子には水平距離 -1、右の子には +1 を設定してキューに追加する。
4. 走査後、map の各キー(垂直列)ごとに合計を計算し、最大値を求める。

C++での実装例

#include<iostream>
#include<map>
#include<vector>
#include<queue>
using namespace std;
class Node {
   public:
   int key;
   Node *left, *right;
};
Node* getNode(int key){
   Node* node = new Node;
   node->key = key;
   node->left = node->right = NULL;
   return node;
}
int getMaxSum(Node* root) {
   if (!root)
   return -1;
   int n = 0;
   int k_node = -1;
   map<int, vector<int> > current_map;
   int hd = 0;
   queue<pair<Node*, int> > que;
   que.push(make_pair(root, hd));
   while (!que.empty()) {
      pair<Node*, int> temp = que.front();
      que.pop();
      hd = temp.second;
      Node* node = temp.first;
      current_map[hd].push_back(node->key);
      if (node->left != NULL)
      que.push(make_pair(node->left, hd - 1));
      if (node->right != NULL)
      que.push(make_pair(node->right, hd + 1));
   }
   map<int, vector<int> >::iterator it;
   int maximum = INT_MIN;
   for (it = current_map.begin(); it != current_map.end(); it++) {
      int temp_sum = 0;
      for (int i = 0; i < it->second.size(); ++i) {
         temp_sum += it->second[i];
      }
      if(maximum < temp_sum){
         maximum = temp_sum;
      }
   }
   return maximum;
}
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);
   root->right->right->right = getNode(9);
   cout << "Maximum sum of vertical nodes: " << getMaxSum(root);
}

出力

Maximum sum of vertical nodes: 12

計算量

ノード数を n とすると、map への挿入や走査に O(log n) のコストがかかるため、時間計算量は O(n log n)、空間計算量は O(n) となります。木が空の場合は -1 を返すことで安全に処理できます。