C++で二分探索木の中順後続ノード(Inorder Successor)を求めるプログラム

二分探索木(BST)とあるノードの値が与えられたとき、そのノードの「中順後続ノード(Inorder Successor)」を求めることを考えます。中順後続ノードとは、ノード p の値よりも大きいキーの中で、最小の値を持つノードのことです。

例として、次のような二分探索木を考えてみましょう。

このとき p = 1 とすると、1 より大きい値の中で最小のものは 2 なので、出力は 2 になります。

解法のアプローチ

この問題は、二分探索木の性質(左の子 < 親 < 右の子)を利用すると、再帰的に効率よく解けます。手順は以下の通りです。

• 再帰メソッド inorderSuccessor(root, p) を定義します。
• root が NULL の場合:
  • NULL を返します。
• root の値が p の値以下の場合:
  • 答えは必ず右部分木に存在するため、inorderSuccessor(root の右の子, p) の結果を返します。
• それ以外の場合(root の値が p の値より大きい場合):
  • root 自身が後続候補となり得るため、まず option = inorderSuccessor(root の左の子, p) を求めます。
  • option が NULL なら root を、NULL でなければ option を返します。

実装例(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;
    }
};
class Solution {
    public:
    TreeNode* inorderSuccessor(TreeNode* root, TreeNode* p) {
        if(!root) return NULL;
        if(root->val <= p->val){
            return inorderSuccessor(root->right, p);
        }else{
            TreeNode* option = inorderSuccessor(root->left, p);
            return !option ? root : option;
        }
    }
};
main(){
    TreeNode *root = new TreeNode(2);
    root->left = new TreeNode(1);
    root->right = new TreeNode(3);
    TreeNode *p = root->left;
    Solution ob;
    cout << (ob.inorderSuccessor(root, p))->val;
}

入力

TreeNode *root = new TreeNode(2);
root->left = new TreeNode(1);
root->right = new TreeNode(3);
1

出力

2

計算量について

このアルゴリズムの時間計算量は O(h) です(h は木の高さ)。各再帰呼び出しで木を1段階ずつ降りていくだけのため、バランスの取れた二分探索木では O(log n)、木が片側に偏っている最悪ケースでは O(n) となります。空間計算量も再帰の深さに依存するため、同様に O(h) です。