C++でBSTの中央値をO(n)時間・O(1)空間で求める方法

基本概念

与えられた二分探索木(BST)に対して、その中央値を求めることが本記事の目的です。

中央値の定義はノード数によって異なります。
・ノード数が奇数の場合：中央値 = (n+1)/2 番目のノードの値
・ノード数が偶数の場合：中央値 = ((n/2 番目のノード + (n+1)/2 番目のノード) / 2

例1：ノード数が奇数のBST

       7
      / \
     4   9
   / \  / \
  2  5  8  10

このBSTを中間順(Inorder)で走査すると「2, 4, 5, 7, 8, 9, 10」というソート済みの列が得られます。したがって、中央値は 7 です。

例2：ノード数が偶数のBST

         7
        / \
       4   9
      / \ /
    2  5 8

このBSTの中間順走査結果は「2, 4, 5, 7, 8, 9」です。
よって、中央値は (5 + 7) / 2 = 6 となります。

アプローチ

BSTの中間順走査結果は常に昇順にソートされているため、中央値を求めるにはまず中間順走査を行う必要があります。この考え方は、「O(1)の追加メモリでBST内のK番目に小さい要素を求める」手法に基づいています。

追加の記憶領域が許可されていれば問題は簡単ですが、再帰やスタックを用いた通常の中間順走査はいずれも空間を消費するため、ここでは使用できません。

そこで採用するのが、余分な空間を一切必要としない「モリス中間順走査(Morris Inorder Traversal)」です。

モリス中間順走査の手順

• current をルートとして初期化します。
• current が NULL でない限り、以下を繰り返します。
  • current に左の子がない場合：
    • current のデータを出力（訪問）します。
    • 右へ移動します（current = current->right）。
  • それ以外の場合：
    • current を、current の左部分木における最も右側のノード（中間順の先行ノード）の右の子として一時的に接続します（スレッドの作成）。
    • 左の子へ移動します（current = current->left）。

この走査では木構造に一時的にスレッド（戻るためのリンク）を作成しますが、処理後に必ず元の構造へ復元されるため、追加メモリは不要です。

最終的な実装方針

1. モリス中間順走査を1回実行し、BSTのノード総数を数えます。
2. 続けて、もう一度モリス中間順走査を実行し、訪問したノードをカウントしながら、カウントが中央値の位置に到達したかどうかを判定します。

ノード数が偶数の場合に対応するため、直前に訪問したノードを指すポインタ（prev）を保持しておきます。

C++による実装例

/* C++ program to find the median of BST in O(n) time and O(1)
space*/
#include<bits/stdc++.h>
using namespace std;
/* Implements a binary search tree Node1 which has data, pointer
to left child and a pointer to right child */
struct Node1{
   int data1;
   struct Node1* left1, *right1;
};
//Shows a utility function to create a new BST node
struct Node1 *newNode(int item1){
   struct Node1 *temp1 = new Node1;
   temp1->data1 = item1;
   temp1->left1 = temp1->right1 = NULL;
   return temp1;
}
/* Shows a utility function to insert a new node with
given key in BST */
struct Node1* insert(struct Node1* node1, int key1){
   /* It has been seen that if the tree is empty, return a new node
   */
   if (node1 == NULL) return newNode(key1);
      /* Else, recur down the tree */
      if (key1 < node1->data1)
         node1->left1 = insert(node1->left1, key1);
      else if (key1 > node1->data1)
         node1->right1 = insert(node1->right1, key1);
         /* return the (unchanged) node pointer */
      return node1;
}
/* Shows function to count nodes in a binary search tree
using Morris Inorder traversal*/
int counNodes(struct Node1 *root1){
   struct Node1 *current1, *pre1;
   // Used to initialise count of nodes as 0
   int count1 = 0;
   if (root1 == NULL)
      return count1;
      current1 = root1;
   while (current1 != NULL){
      if (current1->left1 == NULL){
         // Now count node if its left is NULL
         count1++;
         // Go to its right
         current1 = current1->right1;
      } else {
         /* Determine the inorder predecessor of current */
         pre1 = current1->left1;
         while (pre1->right1 != NULL &&
            pre1->right1 != current1)
            pre1 = pre1->right1;
            /* Construct current1 as right child of its inorder predecessor */
         if(pre1->right1 == NULL){
            pre1->right1 = current1;
            current1 = current1->left1;
         }
         /* we have to revert the changes made in if part to restore the original tree i.e., fix the right child of predecssor */
         else {
            pre1->right1 = NULL;
            // Now increment count if the current
            // node is to be visited
            count1++;
            current1 = current1->right1;
         } /* End of if condition pre1->right1 == NULL */
      } /* End of if condition current1->left1 == NULL*/
   } /* End of while */
   return count1;
}
/* Shows function to find median in O(n) time and O(1) space
using Morris Inorder traversal*/
int findMedian(struct Node1 *root1){
   if (root1 == NULL)
      return 0;
   int count1 = counNodes(root1);
   int currCount1 = 0;
   struct Node1 *current1 = root1, *pre1, *prev1;
   while (current1 != NULL){
      if (current1->left1 == NULL){
         // Now count current node
         currCount1++;
         // Verify if current node is the median
         // Odd case
         if (count1 % 2 != 0 && currCount1 == (count1+1)/2)
            return prev1->data1;
         // Even case
         else if (count1 % 2 == 0 && currCount1 == (count1/2)+1)
            return (prev1->data1 + current1->data1)/2;
            // Now update prev1 for even no. of nodes
         prev1 = current1;
         //Go to the right
         current1 = current1->right1;
      } else {
         /* determine the inorder predecessor of current1 */
         pre1 = current1->left1;
         while (pre1->right1 != NULL && pre1->right1 != current1)
            pre1 = pre1->right1;
         /* Construct current1 as right child of its inorder
         predecessor */
         if (pre1->right1 == NULL){
            pre1->right1 = current1;
            current1 = current1->left1;
         }
         /* We have to revert the changes made in if part to restore the original
         tree i.e., fix the right child of predecssor */
         else {
            pre1->right1 = NULL;
            prev1 = pre1;
            // Now count current node
            currCount1++;
            // Verify if the current node is the median
            if (count1 % 2 != 0 && currCount1 == (count1+1)/2 )
               return current1->data1;
            else if (count1%2==0 && currCount1 == (count1/2)+1)
               return (prev1->data1+current1->data1)/2;
            // Now update prev1 node for the case of even
            // no. of nodes
            prev1 = current1;
            current1 = current1->right1;
         } /* End of if condition pre1->right1 == NULL */
      } /* End of if condition current1->left1 == NULL*/
   } /* End of while */
}
/* Driver program to test above functions*/
int main(){
   /* Let us create following BST
      7
      / \
     4   9
   / \  / \
  2  5 8  10 */
   struct Node1 *root1 = NULL;
   root1 = insert(root1, 7);
   insert(root1, 4);
   insert(root1, 2);
   insert(root1, 5);
   insert(root1, 9);
   insert(root1, 8);
   insert(root1, 10);
   cout << "\nMedian of BST is(for odd no. of nodes) "<< findMedian(root1)         <<endl;
   /* Let us create following BST
       7
      / \
     4   9
    / \  /
   2  5 8
   */
   struct Node1 *root2 = NULL;
   root2 = insert(root2, 7);
   insert(root2, 4);
   insert(root2, 2);
   insert(root2, 5);
   insert(root2, 9);
   insert(root2, 8);
   cout << "\nMedian of BST is(for even no. of nodes) "
   << findMedian(root2);
   return 0;
}

実行結果

Median of BST is(for odd no. of nodes) 7
Median of BST is(for even no. of nodes) 6

計算量について

• 時間計算量： O(n) ― 中間順走査を最大2回行うため、線形時間で完了します。
• 空間計算量： O(1) ― 再帰やスタックを使用せず、ポインタ変数のみで処理するため、定数の追加メモリで済みます。

このように、モリス中間順走査を活用することで、BSTの中央値を追加メモリなしで効率的に求めることができます。メモリ制約の厳しい組み込みシステムや大規模データの処理などで特に有効なテクニックです。