Pythonで二分探索木（BST）の中央値をO(n)時間・O(1)空間で求める方法

問題の概要

二分探索木（Binary Search Tree、BST）が与えられたとき、その中央値を求めることを考えます。ノードの総数を n とすると、中央値は次のように定義されます。

• n が奇数の場合： 中央値 = 中序順（昇順）で (n+1)/2 番目のノードの値
• n が偶数の場合： 中央値 = (n/2 番目のノードの値 + (n+1)/2 番目のノードの値) / 2

例として、次のようなBSTを考えてみましょう。

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

この木の中序走査（昇順）の結果は「2, 4, 5, 7, 8, 9, 10」となります。ノード数は 7 で奇数なので、中央値は (7+1)/2 = 4 番目のノード、つまり 7 です。

アプローチ：モリス順巡回（Morris Traversal）

中序走査を再帰や明示的なスタックで行うと、空間計算量が O(n) にかさんでしまいます。そこで役立つのがモリス順巡回という手法です。各ノードと、中序順における直前のノードとの間に一時的なリンク（スレッド）を張り、走査が済んだら解除することで、スタックなしに中序走査を完結できます。これにより、時間計算量 O(n)・空間計算量 O(1) でBSTの中央値を求められます。

アルゴリズムの手順は以下の通りです。

• root が None の場合は 0 を返す。
• node_count = 木に含まれるノードの総数（こちらもモリス順巡回でカウントする）
• count_curr = 0（ここまでに訪問したノード数）
• current = root
• current が null でない限り、以下を繰り返す。
  • current.left が null の場合：
    • count_curr を 1 増やす。
    • node_count が奇数かつ count_curr = (node_count+1)/2 ならば、previous.data を返す。
    • node_count が偶数かつ count_curr = (node_count/2)+1 ならば、(previous.data + current.data) / 2 を返す。
    • previous = current とし、current = current.right へ進む。
  • それ以外の場合：
    • previous = current.left とし、previous.right が null または current と一致するまで previous = previous.right を辿る（中序順での直前のノードを探す）。
    • previous.right が null の場合（スレッド未作成）：
      • previous.right = current としてスレッドを張り、current = current.left へ進む。
    • そうでない場合（左部分木の走査が完了して戻ってきたとき）：
      • previous.right = None としてスレッドを解除する。
      • count_curr を 1 増やす。
      • node_count が奇数かつ count_curr = (node_count+1)/2 ならば、current.data を返す。
      • node_count が偶数かつ count_curr = (node_count/2)+1 ならば、(previous.data + current.data) / 2 を返す。
      • previous = current とし、current = current.right へ進む。

Pythonでの実装例

理解を深めるために、実際の実装を見てみましょう。

class TreeNode:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None

def number_of_nodes(root):
    node_count = 0
    if (root == None):
        return node_count
    current = root
    while (current != None):
        if (current.left == None):
            node_count += 1
            current = current.right
        else:
            previous = current.left
            while (previous.right != None and previous.right != current):
                previous = previous.right
            if (previous.right == None):
                previous.right = current
                current = current.left
            else:
                previous.right = None
                node_count += 1
                current = current.right
    return node_count

def calculate_median(root):
    if (root == None):
        return 0
    node_count = number_of_nodes(root)
    count_curr = 0
    current = root
    while (current != None):
        if (current.left == None):
            count_curr += 1
            if (node_count % 2 != 0 and count_curr == (node_count + 1) // 2):
                return previous.data
            elif (node_count % 2 == 0 and count_curr == (node_count // 2) + 1):
                return (previous.data + current.data) // 2
            previous = current
            current = current.right
        else:
            previous = current.left
            while (previous.right != None and previous.right != current):
                previous = previous.right
            if (previous.right == None):
                previous.right = current
                current = current.left
            else:
                previous.right = None
                count_curr += 1
                if (node_count % 2 != 0 and count_curr == (node_count + 1) // 2):
                    return current.data
                elif (node_count % 2 == 0 and count_curr == (node_count // 2) + 1):
                    return (previous.data + current.data) // 2
                previous = current
                current = current.right

root = TreeNode(7)
root.left = TreeNode(4)
root.right = TreeNode(9)
root.left.left = TreeNode(2)
root.left.right = TreeNode(5)
root.right.left = TreeNode(8)
root.right.right = TreeNode(10)
print(calculate_median(root))

入力

root = TreeNode(7)
root.left = TreeNode(4)
root.right = TreeNode(9)
root.left.left = TreeNode(2)
root.left.right = TreeNode(5)
root.right.left = TreeNode(8)
root.right.right = TreeNode(10)

出力

7

計算量について

• 時間計算量： O(n)。各ノードを高々2回訪問するだけなので線形時間で処理できます。
• 空間計算量： O(1)。再帰呼び出しやスタックを一切使わず、ポインタ操作のみで完結します。

なお、サンプルコードでは整数除算（//）を使用しているため、ノード数が偶数の場合の中央値は小数点以下が切り捨てられた整数として返される点に注意してください。