C++で二分探索木（BST）から指定した合計値になるペアをすべて検索する方法

このチュートリアルでは、二分探索木（BST）の中から、合計が指定された数値と等しくなるすべてのペアを見つけるプログラムを作成します。

ペアを効率よく探すために、木のノードの値を2つの異なるリストに格納しながら処理を進めます。これは、ソート済み配列で使われる「Two Pointer（双方向ポインタ）」のテクニックを、BSTに対して左右両方向から中順走査を行うことで実現するアプローチです。それでは、問題を解くための手順を順番に見ていきましょう。

アルゴリズムの手順

• 二分木用の構造体（struct Node）を作成します。

• 新しいノードを二分探索木に挿入する関数を用意します。
  ※二分探索木では、根（ルート）より小さい要素は必ず左側に、大きい要素は右側に配置されるという性質があります。

• 木の左側ノードと右側ノードをそれぞれ格納するための空のリストを2つ初期化します。

• 左または右のノードがNULLになるまで、二分探索木を反復処理します。

  • ループを使って、左方向のすべての要素を左側ノードのリストに格納します。

  • 同様に、右方向のすべての要素を右側ノードのリストに格納します。

  • 各リストから末尾のノードを取得します。

  • 取得した2つの値を比較します。

    • 左側のノードの値が右側のノードの値以上になった場合は、ループを抜けます。

    • 2つの値の合計が指定された数値と等しい場合は、そのペアを出力し、両方のリストから該当ノードを削除します。

    • 合計が指定された数値より小さい場合は、左側リストから末尾のノードを取り除き、そのノードの右部分木へ移動します。

    • 合計が指定された数値より大きい場合は、右側リストから末尾のノードを取り除き、そのノードの左部分木へ移動します。

実装例

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

#include<bits/stdc++.h>
using namespace std;
struct Node {
    int data;
    Node *left, *right, *root;
    Node(int data) {
        this->data = data;
        left = NULL;
        right = NULL;
        root = NULL;
    }
};
Node* insertNewNode(Node *root, int data) {
    if (root == NULL) {
        root = new Node(data);
        return root;
    }
    if (root->data < data) {
        root->right = insertNewNode(root->right, data);
    }
    else if (root->data > data) {
        root->left = insertNewNode(root->left, data);
    }
    return root;
}
void findThePairs(Node *node, int target) {
    vector<Node*> left_side_nodes;
    vector<Node*> right_side_nodes;
    Node *current_left = node;
    Node *current_right = node;
    while (current_left != NULL || current_right != NULL || (left_side_nodes.size() > 0 && right_side_nodes.size() > 0)) {
        while (current_left != NULL) {
            left_side_nodes.push_back(current_left);
            current_left = current_left->left;
        }
        while (current_right != NULL) {
            right_side_nodes.push_back(current_right);
            current_right = current_right->right;
        }
        Node *left_side_node = left_side_nodes[left_side_nodes.size() - 1];
        Node *right_side_node = right_side_nodes[right_side_nodes.size() - 1];
        int left_side_value = left_side_node->data;
        int right_side_value = right_side_node->data;
        if (left_side_value >= right_side_value) {
            break;
        }
        if (left_side_value + right_side_value < target) {
            left_side_nodes.pop_back();
            current_left = left_side_node->right;
        }
        else if (left_side_value + right_side_value > target) {
            right_side_nodes.pop_back();
            current_right = right_side_node->left;
        }
        else {
            cout << left_side_node->data << " " << right_side_node->data << endl;
            right_side_nodes.pop_back();
            left_side_nodes.pop_back();
            current_left = left_side_node->right;
            current_right = right_side_node->left;
        }
    }
}
int main() {
    Node *root = NULL;
    root = insertNewNode(root, 25);
    root = insertNewNode(root, 20);
    root = insertNewNode(root, 30);
    root = insertNewNode(root, 15);
    root = insertNewNode(root, 21);
    root = insertNewNode(root, 19);
    root = insertNewNode(root, 31);
    findThePairs(root, 50);
}

出力結果

上記のコードを実行すると、次のような結果が得られます。

19 31
20 30

まとめ

このアルゴリズムは、BSTを左端（最小値）と右端（最大値）から同時に走査することで、時間計算量O(n)、空間計算量O(h)（hは木の高さ）で目的のペアをすべて検索できます。配列のTwo Pointer手法を木構造に応用した実装として、非常に有用なパターンなのでぜひ覚えておきましょう。

チュートリアルの内容についてご不明な点がある場合は、コメント欄でお気軽にお知らせください。