C++で二分探索木(BST)を復元するアルゴリズムを解説

問題の概要

二分探索木(BST)において、誤って2つのノードの値が入れ替わってしまった状態を考えます。この記事では、入れ替わった2つのノードを特定し、木を正しい二分探索木の状態へ復元する方法をC++で解説します。

例えば、次のような木が与えられた場合(左図)、復元後の木は右図のようになります。

解決のアプローチ

二分探索木には「中順(インオーダー)走査を行うと、ノードの値が昇順に並ぶ」という重要な性質があります。この性質を利用すると、入れ替わったノードを効率的に検出できます。具体的には、以下の手順で解決します。

1. prev、first、second という3つのノード参照を用意します。
2. findProblem() というメソッドを定義します(引数としてノードを受け取ります)。
3. ノードがNULLの場合は何もせずに戻ります。
4. findProblem(ノードの左部分木)を再帰的に呼び出します。
5. prev がNULLでなく、かつ prev の値が現在のノードの値より大きい場合(昇順が崩れている箇所を検出):
   • first がNULLであれば、first = prev とします。
   • second = 現在のノード とします。
6. prev = 現在のノード として更新します。
7. findProblem(ノードの右部分木)を再帰的に呼び出します。
8. メインの recoverTree() メソッドでは、prev・first・second をNULLで初期化し、findProblem(root) を呼び出した後、first と second の値を交換します。

なぜこの方法で動くのか

中順走査ではBSTの値が昇順に並ぶはずです。2つの要素が入れ替わると、昇順が崩れる箇所が1〜2か所現れます。最初に検出された違反箇所の「前のノード」が first に、「現在のノード」が second に記録され、最後にこの2つの値を交換することで木が復元されます。

C++での実装例

以下の実装例を見ると、より理解が深まるでしょう。

#include <bits/stdc++.h>
using namespace std;
class TreeNode{
    public:
    int val;
    TreeNode *left, *right;
    TreeNode(int data){
        val = data;
        left = NULL;
        right = NULL;
    }
};
void insert(TreeNode **root, int val){
    queue<TreeNode*> q;
    q.push(*root);
    while(q.size()){
        TreeNode *temp = q.front();
        q.pop();
        if(!temp->left){
            if(val != NULL)
                temp->left = new TreeNode(val);
            else
                temp->left = new TreeNode(0);
            return;
        } else {
            q.push(temp->left);
        }
        if(!temp->right){
            if(val != NULL)
                temp->right = new TreeNode(val);
            else
                temp->right = new TreeNode(0);
            return;
        } else {
            q.push(temp->right);
        }
    }
}
TreeNode *make_tree(vector<int> v){
    TreeNode *root = new TreeNode(v[0]);
    for(int i = 1; i<v.size(); i++){
        insert(&root, v[i]);
    }
    return root;
}
void tree_level_trav(TreeNode* root){
    if (root == NULL) return;
    cout << "[";
    queue<TreeNode*> q;
    TreeNode *curr;
    q.push(root);
    q.push(NULL);
    while (q.size() > 1) {
        curr = q.front();
        q.pop();
        if (curr == NULL){
            q.push(NULL);
        } else {
            if(curr->left)
                q.push(curr->left);
            if(curr->right)
                q.push(curr->right);
            if(curr->val == 0){
                cout << "null" << ", ";
            }else{
                cout << curr->val << ", ";
            }
        }
    }
    cout << "]"<<endl;
}
class Solution {
    public:
    TreeNode* prev;
    TreeNode* first;
    TreeNode* second;
    void swapValue(TreeNode* first, TreeNode* second){
        int x = first->val;
        first->val = second -> val;
        second->val = x;
    }
    void findProblem(TreeNode* node){
        if(!node || node->val == 0)return;
        findProblem(node->left);
        if((prev!=NULL && prev->val != 0) && prev->val > node->val){
            if(!first){
                first = prev;
            }
            second = node;
        }
        prev = node;
        findProblem(node->right);
    }
    void recoverTree(TreeNode* root) {
        prev = first = second = NULL;
        findProblem(root);
        swapValue(first, second);
    }
};
int main(){
    vector<int> v = {1,3,NULL,NULL,2};
    TreeNode *root = make_tree(v);
    Solution ob;
    ob.recoverTree(root);
    tree_level_trav(root);
}

入力

{1,3,NULL,NULL,2}

出力

[3, 1, null, null, 2]

このように、中順走査の性質を活かすことで、木全体を再構築することなく、入れ替わった2つのノードの値だけを交換してBSTを復元できます。計算量はノード数をNとすると O(N)、再帰の深さが木の高さに依存するため空間計算量は O(H)(Hは木の高さ)となります。