C++で2つのBSTから合計が指定値xと等しいペアを数える方法

2つの二分探索木（BST）と整数値 x が与えられます。この記事の目的は、BST_1 から1つのノード、BST_2 からもう1つのノードを選んだペアのうち、両ノードの値の合計が x に一致するものの個数を求めることです。具体的には、BST_1 のノードと BST_2 のノードのデータ部分を加算し、その合計が x と等しければカウントを1つ増やしていきます。

具体例で確認してみましょう。

入力

出力 − 合計が指定値 x に等しい2つのBSTからのペアの数 − 1

説明 − 該当するペアは (8, 6) です。

入力

出力 − 合計が指定値 x に等しい2つのBSTからのペアの数 − 2

説明 − 該当するペアは (5, 15) と (4, 16) です。

プログラムで使用するアプローチ

このアプローチでは、スタックを利用した反復的な中順走査（インオーダー走査）によって2つのBSTを同時に走査します。BST_1 は最小ノードから最大ノードへ向かう昇順の中順走査で辿り、BST_2 は逆順（降順）の中順走査で辿ります。そして、両方のBSTにおける現在ノードの値の合計を評価します。合計が x と等しければカウントを増やし、合計が x より大きければ BST_2 側を次に小さいノード（先行ノード）へ進め、合計が x より小さければ BST_1 側を次に大きいノード（後続ノード）へ進めます。これは、ソート済み配列から合計が特定の値になるペアを探す「双方向ポインタ」手法を、2つのBSTに応用したものと言えます。

• 整数型のデータ部分と、子ノードを指す left・right ポインタを持つ2つの木 BST_1 と BST_2 を用意します。
• 関数 insert_node(int data) は、指定したデータを持つ新しいノードを作成し、そのポインタを返します。
• insert_node() を使って両方のBSTを構築し、BST_sum_x(Tree* BST_1, Tree* BST_2, int x) に渡します。
• 関数 BST_sum_x(...) は、両方の木の根ノードと目標値 x を受け取り、データ部分の合計が x になるノードのペアの個数を返します。
• 合計が x となるペアの数を数えるため、カウントの初期値を 0 とします。
• 反復的な中順走査のために、Tree* 型の変数 stack_top_1 と stack_top_2 を宣言します。
• 2つのスタック stack_1 と stack_2 を作成します。
• 外側の while ループを開始します。
• while ループで BST_1 の最左端（最小値）ノードまで進み、経路上のすべてのノードを stack_1 にプッシュします。
• while ループで BST_2 の最右端（最大値）ノードまで進み、経路上のすべてのノードを stack_2 にプッシュします。
• どちらかのスタックが空になったら、外側の while ループを抜けます。
• 両スタックのトップにあるノードのデータ部分を加算し、temp に格納します。
• temp（合計）== x の場合は count をインクリメントし、pop 操作で stack_1 と stack_2 の両方からトップ要素を削除します。その後、BST_1 = stack_top_1->right、BST_2 = stack_top_2->left と設定します（それぞれ BST_1 の次の後続ノード、BST_2 の次の先行ノード）。
• temp < x の場合は stack_1 のみからトップを削除し、BST_1 の次の後続ノードへ移動します。
• temp > x の場合は stack_2 のみからトップを削除し、BST_2 の次の先行ノードへ移動します。
• 外側の while ループが終了した時点で、count には合計が x となるペアの総数が格納されています。
• count を結果として返します。

なお、この手法の時間計算量は O(n1 + n2)、空間計算量は O(h1 + h2) となります（n1・n2 は各木のノード数、h1・h2 は各木の高さ）。すべてのノードの組み合わせを総当たりで調べる O(n1 × n2) の素朴な方法に比べ、大幅に効率的である点が大きな利点です。

例

#include <bits/stdc++.h>
using namespace std;
struct Tree{
    int data;
    Tree* left, *right;
};
Tree* insert_node(int data){
    Tree* newNode = (Tree*)malloc(sizeof(Tree));
    newNode->data = data;
    newNode->left = NULL;
    newNode->right = NULL;
}
int BST_sum_x(Tree* BST_1, Tree* BST_2, int x){
    int count = 0;
    Tree* stack_top_1, *stack_top_2;
    stack<Tree*> stack_1, stack_2;
    if (BST_1 == NULL || BST_2 == NULL){
        return 0;
    }
    while (1){
        while (BST_1 != NULL){
            stack_1.push(BST_1);
            BST_1 = BST_1->left;
        }
        while (BST_2 != NULL){
            stack_2.push(BST_2);
            BST_2 = BST_2->right;
        }
        if (stack_1.empty() || stack_2.empty()){
            break;
        }
        stack_top_1 = stack_1.top();
        stack_top_2 = stack_2.top();
        int temp = stack_top_1->data + stack_top_2->data;
        if (temp == x){
            count++;
            stack_1.pop();
            stack_2.pop();
            BST_1 = stack_top_1->right;
            BST_2 = stack_top_2->left;
        }
        else if (temp < x){
            stack_1.pop();
            BST_1 = stack_top_1->right;
        }
        else{
            stack_2.pop();
            BST_2 = stack_top_2->left;
        }
    }
    return count;
}
int main(){
    //BST 1
    Tree* BST_1 = insert_node(15);
    BST_1->left = insert_node(10);
    BST_1->right = insert_node(8);
    BST_1->left->left = insert_node(12);
    BST_1->left->right = insert_node(24);
    BST_1->right->left = insert_node(16);
    //BST 2
    Tree* BST_2 = insert_node(20);
    BST_2->left = insert_node(16);
    BST_2->right = insert_node(4);
    BST_2->left->left = insert_node(18);
    BST_2->left->right = insert_node(28);
    BST_2->right->left = insert_node(22);
    int x = 28;
    cout<<"Count of pairs from two BSTs whose sum is equal to a given value x ar: "<<BST_sum_x(BST_1, BST_2, x);
    return 0;
}

出力

上記のコードを実行すると、次の出力が生成されます −

Count of pairs from two BSTs whose sum is equal to a given value x ar: 1