Pythonで理解する蟻本「2-4 二分探索木」(p.75)

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 二分探索木」(p.75)
のコードをPythonで書き直したものとなっています。

二分探索木の実装

コード

# ノードを表す構造体
class node:
    def __init__(self):
        self.val = None
        self.lch = None
        self.rch = None

# 数xを追加
def insert(p, x):
    if p == None:
        q = node()
        q.val = x
        return q
    else:
        if x < p.val:
            p.lch = insert(p.lch, x)
        else:
            p.rch = insert(p.rch, x)
        return p

# 数xを検索
def find(p, x):
    if p == None:
        return False
    elif x == p.val:
        return True
    elif x < p.val:
        return find(p.lch, x)
    else:
        return find(p.rch, x)

# 数xを削除
def remove(p, x):
    if p == None:
        return None
    elif x < p.val:
        p.lch = remove(p.lch, x)
    elif x > p.val:
        p.rch = remove(p.rch, x)
    elif p.lch == None:
        q = p.rch
        del p
        return q
    elif p.lch.rch == None:
        q = p.lch
        q.rch = p.rch
        del p
        return q
    else:
        q = node()
        q = p.lch
        while q.rch.rch != None:
            q = q.rch
        r = q.rch
        q.rch = r.lch
        r.lch = p.lch
        r.rch = p.rch
        del p
        return r
    return p

実行例

コード

root = None
root = insert(root, 1)    # 1を追加
print(find(root, 1))      # 1を検索
root = remove(root, 1)    # 1を削除
print(find(root, 1))      # 1を検索

実行結果

True
False

プログラミング言語の標準ライブラリ

残念ながらPythonにはC++のsetやmapに相当するコンテナは存在しません。

類似した物としてはbisectモジュールが挙げられます。

bisectモジュールを使えば二分探索によってO(log n)で要素を検索する事は可能です。

しかし、要素を挿入する際にO(n)かかってしまうため、挿入の際の全体の計算量はO(n)となってしまいます。