JavaのTreesetをPythonで

「できらあ！！」

「え！！同じデータ構造をPythonで！？」

というわけでJavaのTreesetに当たるデータ構造がPythonで欲しくなったので、自分で実装してみました。

treeset.pyをコピーして使ってください。

Treesetはsetのように、重複が存在しないデータ集合ですが、setと違い、初期時にデータがソートされます。また、新しい要素が追加されるたびに、二分木探索によって適切な位置に挿入されます。

使い方は例えば、こんな感じ。

from treeset import TreeSet

# 初期化時に重複要素は消える、またソートが行われる
ts = TreeSet([3,7,2,7,1,3])
print(ts)
>>> [1, 2, 3, 7]

# 4を加えると、3と7の間に挿入される
ts.add(4)
print(ts)
>>> [1, 2, 3, 4, 7]

インデックスによる取得、要素を指定してリムーブもできます。

# 三番目の要素を取得
print(ts[2])
>>> 3

# 最後の要素を取得
print(ts[2])
>>> 3

# Treesetに7があれば消去
ts.remove(7) print(ts) >>> [1, 2, 3, 4] ts.pop(2)
# インデックスをpop print(ts) >>> [1, 2, 4]

中身はシンプルでbisectモジュールを使って、適宜、中身のデータをソートしているだけです。

from treeset import TreeSet

class TreeSet(object)

    def __init__(self, elements):
        self._treeset = []
        for element in elements:
            self.add(element)

    def add(self, element):
        # 要素が既に存在しなければbisect.insertで適切な位置に挿入
        if element in self._treeset: continue
        self.add(element)
        bisect.insort(self._treeset, element)

よかったらお試しください。