pythonで平衡二分木を使った辞書を作ってみた（C++ std::map相当)

# はじめに
pythonのset（集合）やdict（辞書）は値の検索削除がO(1)と高速ですが、要素の大きさに順序を持つことができず、順序を気にしながらの実装を愚直にすると計算時間が大きくかかってしまうなどの課題があります。一方でc++では、順序付き集合のsetや順序付きの連想配列のmap（辞書相当）が標準で整備されており、競技プログラミングなど時間制約がタイトな場合では、標準ライブラリーの差でもpython勢は不利になることが多いです。そこで今回はc++のmapに相当する平衡二分木を用いた辞書の実装例を紹介したいと思います。

# 参考記事
今回は平衡二分木部分に関しては出来合いのものを用いました。
<a href="https://stnkien.hatenablog.com/entry/avl-tree">Pythonで非再帰AVL木 - 競プロ記録 他</a>
(非再帰での実装となっておりpythonでも高速に動くためとても重宝しております）
こちらで紹介されている記事の実装例にいくつか機能を付与する形で実装しました。

# コードの紹介
コードがとても長くなってしまっているので、関数及び使用方法について次項で重点的に説明しようと思います。
<details><summary><font color="red"><strong>コードの表示はこちらをクリック</strong></font></summary><div>

```AVLTree_dict.py
#非再帰の平衡二分木
import copy
class Node:
    """ノード

    Attributes:
        key (any): ノードのキー。比較可能なものであれば良い。(1, 4)などタプルも可。
        val (any): ノードの値。
        lch (Node): 左の子ノード。
        rch (Node): 右の子ノード。
        bias (int): 平衡度。(左部分木の高さ)-(右部分木の高さ)。
        size (int): 自分を根とする部分木の大きさ

    """

    def __init__(self, key, val):
        self.key = key
        self.val = val
        self.lch = None
        self.rch = None
        self.bias = 0
        self.size = 1

class AVLTree:
    """非再帰AVL木

    Attributes:
        root (Node): 根ノード。初期値はNone。
        valdefault (any): ノード値のデフォルト値。デフォルトではNone。（数値、リストなど可）

    """

    def __init__(self,valdefault=None):
        self.root = None
        self.valdefault = valdefault
        
    def rotate_left(self, v):
        u = v.rch
        u.size = v.size
        v.size -= u.rch.size + 1 if u.rch is not None else 1
        v.rch = u.lch
        u.lch = v
        if u.bias == -1:
            u.bias = v.bias = 0
        else:
            u.bias = 1
            v.bias = -1
        return u

    def rotate_right(self, v):
        u = v.lch
        u.size = v.size
        v.size -= u.lch.size + 1 if u.lch is not None else 1
        v.lch = u.rch
        u.rch = v
        if u.bias == 1:
            u.bias = v.bias = 0
        else:
            u.bias = -1
            v.bias = 1
        return u

    def rotateLR(self, v):
        u = v.lch
        t = u.rch
        t.size = v.size
        v.size -= u.size - (t.rch.size if t.rch is not None else 0)
        u.size -= t.rch.size + 1 if t.rch is not None else 1
        u.rch = t.lch
        t.lch = u
        v.lch = t.rch
        t.rch = v
        self.update_bias_double(t)
        return t

    def rotateRL(self, v):
        u = v.rch
        t = u.lch
        t.size = v.size
        v.size -= u.size - (t.lch.size if t.lch is not None else 0)
        u.size -= t.lch.size + 1 if t.lch is not None else 1
        u.lch = t.rch
        t.rch = u
        v.rch = t.lch
        t.lch = v
        self.update_bias_double(t)
        return t

    def update_bias_double(self, v):
        if v.bias == 1:
            v.rch.bias = -1
            v.lch.bias = 0
        elif v.bias == -1:
            v.rch.bias = 0
            v.lch.bias = 1
        else:
            v.rch.bias = 0
            v.lch.bias = 0
        v.bias = 0

    def insert(self, key, val=None):
        """挿入

        指定したkeyを挿入する。valはkeyのノード値。

        Args:
            key (any): キー。
            val (any): 値。（指定しない場合はvaldefaultが入る）

        Note:
            同じキーがあった場合は上書きする。

        """
        if val == None:
            val = copy.deepcopy(self.valdefault)
        if self.root is None:
            self.root = Node(key, val)
            return

        v = self.root
        history = []
        while v is not None:
            if key < v.key:
                history.append((v, 1))
                v = v.lch
            elif v.key < key:
                history.append((v, -1))
                v = v.rch
            elif v.key == key:
                v.val = val
                return

        p, pdir = history[-1]
        if pdir == 1:
            p.lch = Node(key, val)
        else:
            p.rch = Node(key, val)

        while history:
            v, direction = history.pop()
            v.bias += direction
            v.size += 1

            new_v = None
            b = v.bias
            if b == 0:
                break

            if b == 2:
                u = v.lch
                if u.bias == -1:
                    new_v = self.rotateLR(v)
                else:
                    new_v = self.rotate_right(v)
                break
            if b == -2:
                u = v.rch
                if u.bias == 1:
                    new_v = self.rotateRL(v)
                else:
                    new_v = self.rotate_left(v)
                break

        if new_v is not None:
            if len(history) == 0:
                self.root = new_v
                return
            p, pdir = history.pop()
            p.size += 1
            if pdir == 1:
                p.lch = new_v
            else:
                p.rch = new_v

        while history:
            p, pdir = history.pop()
            p.size += 1

    def delete(self, key):
        """削除

        指定したkeyの要素を削除する。

        Args:
            key (any): キー。

        Return:
            bool: 指定したキーが存在するならTrue、しないならFalse。

        """
        v = self.root
        history = []
        while v is not None:
            if key < v.key:
                history.append((v, 1))
                v = v.lch
            elif v.key < key:
                history.append((v, -1))
                v = v.rch
            else:
                break
        else:
            return False

        if v.lch is not None:
            history.append((v, 1))
            lmax = v.lch
            while lmax.rch is not None:
                history.append((lmax, -1))
                lmax = lmax.rch

            v.key = lmax.key
            v.val = lmax.val

            v = lmax

        c = v.rch if v.lch is None else v.lch

        if history:
            p, pdir = history[-1]
            if pdir == 1:
                p.lch = c
            else:
                p.rch = c
        else:
            self.root = c
            return True

        while history:
            new_p = None

            p, pdir = history.pop()
            p.bias -= pdir
            p.size -= 1

            b = p.bias
            if b == 2:
                if p.lch.bias == -1:
                    new_p = self.rotateLR(p)
                else:
                    new_p = self.rotate_right(p)

            elif b == -2:
                if p.rch.bias == 1:
                    new_p = self.rotateRL(p)
                else:
                    new_p = self.rotate_left(p)

            elif b != 0:
                break

            if new_p is not None:
                if len(history) == 0:
                    self.root = new_p
                    return True
                gp, gpdir = history[-1]
                if gpdir == 1:
                    gp.lch = new_p
                else:
                    gp.rch = new_p
                if new_p.bias != 0:
                    break

        while history:
            p, pdir = history.pop()
            p.size -= 1

        return True

    def member(self, key):
        """キーの存在チェック

        指定したkeyがあるかどうか判定する。

        Args:
            key (any): キー。

        Return:
            bool: 指定したキーが存在するならTrue、しないならFalse。

        """
        v = self.root
        while v is not None:
            if key < v.key:
                v = v.lch
            elif v.key < key:
                v = v.rch
            else:
                return True
        return False

    def getval(self, key):
        """値の取り出し

        指定したkeyの値を返す。

        Args:
            key (any): キー。

        Return:
            any: 指定したキーが存在するならそのオブジェクト。存在しなければvaldefault

        """
        v = self.root
        while v is not None:
            if key < v.key:
                v = v.lch
            elif v.key < key:
                v = v.rch
            else:
                return v.val
        self.insert(key)
        return self[key] #

    def lower_bound(self, key):
        """下限つき探索

        指定したkey以上で最小のキーを見つける。[key,inf)で最小

        Args:
            key (any): キーの下限。

        Return:
            any: 条件を満たすようなキー。そのようなキーが一つも存在しないならNone。

        """
        ret = None
        v = self.root
        while v is not None:
            if v.key >= key:
                if ret is None or ret > v.key:
                    ret = v.key
                v = v.lch
            else:
                v = v.rch
        return ret

    def upper_bound(self, key):
        """上限つき探索

        指定したkey未満で最大のキーを見つける。[-inf,key)で最大

        Args:
            key (any): キーの上限。

        Return:
            any: 条件を満たすようなキー。そのようなキーが一つも存在しないならNone。

        """
        ret = None
        v = self.root
        while v is not None:
            if v.key < key:
                if ret is None or ret < v.key:
                    ret = v.key
                v = v.rch
            else:
                v = v.lch
        return ret

    def find_kth_element(self, k):
        """小さい方からk番目の要素を見つける

        Args:
            k (int): 何番目の要素か(0オリジン)。

        Return:
            any: 小さい方からk番目のキーの値。
        """
        v = self.root
        s = 0
        while v is not None:
            t = s+v.lch.size if v.lch is not None else s
            if t == k:
                return v.key
            elif t < k:
                s = t+1
                v = v.rch
            else:
                v = v.lch
        return None
    
    def getmin(self):
        '''
        Return:
            any: 存在するキーの最小値
        '''
        if len(self) == 0:
            raise('empty')
        ret = None
        v = self.root
        while True:
            ret = v
            v = v.lch
            if v == None:
                break
        return ret.key
    
    def getmax(self):
        '''
        Return:
            any: 存在するキーの最大値
        '''
        if len(self) == 0:
            raise('empty')
        ret = None
        v = self.root
        while True:
            ret = v
            v = v.rch
            if v == None:
                break
        return ret.key
    
    def popmin(self):
        '''
        存在するキーの最小値をpopする
        Return:
            any: popした値
        '''
        if len(self) == 0:
            raise('empty')
        ret = None
        v = self.root
        while True:
            ret = v
            v = v.lch
            if v == None:
                break      
        del self[ret.key]
        return ret.key
    
    def popmax(self):
        '''
        存在するキーの最大値をpopする
        Return:
            any: popした値
        '''
        if len(self) == 0:
            raise('empty')
        ret = None
        v = self.root
        while True:
            ret = v
            v = v.rch
            if v == None:
                break
        del self[ret.key]
        return ret.key

    def popkth(self,k):
        '''
        存在するキーの小さい方からk番目をpopする
        Return:
            any: popした値
        '''
        key = self.find_kth_element(k)
        del self[key]
        return key
    
    def get_key_val(self):
        '''
        Return:
            dict: 存在するキーとノード値をdictで出力
        '''
        retdict = dict()
        for i in range(len(self)):
            key = self.find_kth_element(i)
            val = self[key]
            retdict[key] = val
        return retdict
    
    def values(self):
        for i in range(len(self)):
            yield self[self.find_kth_element(i)]

    def keys(self):
        for i in range(len(self)):
            yield self.find_kth_element(i)
    
    def items(self):
        for i in range(len(self)):
            key = self.find_kth_element(i)
            yield key,self[key]

    def __iter__(self): return self.keys()
    def __contains__(self, key): return self.member(key)
    def __getitem__(self, key): return self.getval(key)
    def __setitem__(self, key, val): return self.insert(key, val)
    def __delitem__(self, key): return self.delete(key)
    def __bool__(self): return self.root is not None
    def __len__(self): return self.root.size if self.root is not None else 0
    def __str__(self): return str(type(self))+'('+str(self.get_key_val())+')'
```

</div></details>

# 関数の説明

|関数  |説明  |
|---|---|
|<strong>rotate_left - update_bias_double</strong> |平衡二分木を実現するうえで最も大事な部分です（平衡を保つための回転などを行っています）。外部で呼び出す必要はないので、使用上は気にしなくていい部分です。  |
|<strong>insert(key,val=None)</strong>|keyを平衡二分木に挿入します。valはkeyのノード値(value)です（辞書のkey,valueと同等)|
|<strong>delete(key)</strong>|指定したkeyの要素を削除します。|
|<strong>menber(key)</strong>|指定したkeyの有無を判定します。|
|<strong>getval(key)</strong>|指定したkeyの値のノード値を返します。(dict[key]と同等）|
|<strong>lower_bound(key)</strong>|[key,inf)で最小のものを返します。|
|<strong>upper_bound(key)</strong>|[-inf,key)で最大のものを返します。|
|<strong>find_kth_element(k)</strong>|小さい方からk番目のkeyを返します（最小を0番目とする）|


以下元記事にはなかった追加機能です。

|関数  |説明  |
|---|---|
|<strong>getmin()</strong>|最小のkeyを返します。|
|<strong>getmax()</strong>|最大のkeyを返します。|
|<strong>popmin()</strong>|最小のkeyを返し、削除します。(heapqと同じ操作が実現可能）|
|<strong>popmax()</strong>|最大のkeyを返し、削除します。|
|<strong>popkth(k)</strong>|小さい方からk番目のkeyを返し、削除します。|
|<strong>get_key_val()</strong>|存在するすべてのkeyとノード値(value)をdictで返します。|
|<strong>values()</strong>|すべてのノード値(value)を返します。(dictのvalues()と同等)|
|<strong>keys()</strong>|すべてのkeyを返します。(dictのkeys()と同等)|
|<strong>items()</strong>|すべてのkeyとノード値(value)を返します。(dictのitems()と同等)|

計算量はget_key_val(),values(),keys(),items()がO(Nlog(N))、他はO(log(N))と高速です。
その他に、使用感をdictに合わせるために特殊メゾットを定義しています。

# 使用例
注意点としてはkeyは大小比較可能なのものである必要があるというところです（数値、文字列、tupleなどは可能）

### ・普通の平衡二分木としての動作

```a.py
AVL = AVLTree()
AVL.insert(10)
AVL.insert(20)
AVL.insert(30)
AVL.insert(40)
AVL.insert(50)
print(AVL.lower_bound(15))
# 20
print(AVL.find_kth_element(2))
# 30
print(40 in AVL)
# True
del AVL[40]   # AVL.delete(40)と等価
print(40 in AVL)
# False
print(list(AVL))
# [10, 20, 30, 50]
print(AVL.popmin())
# 10
print(AVL.popkth(1)) # 20,30,50のうち1番目(0オリジン)の30
# 30
print(list(AVL))
# [20, 50]
print(len(AVL))
# 2
```

### ・辞書っぽい使い方
```b.py
AVL1 = AVLTree()
AVL1['a'] = 'A'
AVL1['b'] = 'B'
AVL1['f'] = 'C'
AVL1['aa'] = 'AA'
print(list(AVL1))
# ['a', 'aa', 'b', 'f']
print(AVL1.get_key_val())
# {'a': 'A', 'aa': 'AA', 'b': 'B', 'f': 'C'}
print(AVL1.getmax())
# f
print(AVL1.upper_bound('e'))
# b
```

### ・collections.defaultdict相当の動作

```c.py
# はじめにvaldefaultを指定することでdefaultdictに相当する処理が可能。
AVL2 = AVLTree(valdefault=[])
AVL2[20].append(2)
AVL2[20].append(3)
AVL2[20].append(6)
AVL2[30].append(5)
AVL2[40].append(1)
AVL2[40].append(2)
print(AVL2.get_key_val())
# {20: [2, 3, 6], 30: [5], 40: [1, 2]}
print(AVL2[20].pop())
# 6
print(40 in AVL2)
# True
print(50 in AVL2)
# False
print(AVL2.popmax())
# 40
AVL2[50].append(5)
AVL2[50].append(6)
print(AVL2.get_key_val())
# {20: [2, 3], 30: [5], 50: [5, 6]}
```

### ・collections.Counter相当の動作

```d.py
AVL3 = AVLTree(valdefault=0)
AVL3[30] += 3
AVL3[40] += 2
AVL3[50] += 1
print(AVL3.get_key_val())
# {30: 3, 40: 2, 50: 1}
AVL3[50] += 5
print(AVL3.get_key_val())
# {30: 3, 40: 2, 50: 6}
print(list(AVL3.values()))
# [3, 2, 6]
for key,val in AVL3.items():
    print(key,val)
    # 30 3
    # 40 2
    # 50 6
while AVL3:
    key = AVL3.getmin()
    print(key,AVL3[key])
    AVL3.popmin()
    # 30 3
    # 40 2
    # 50 6
```

# 競技プログラミングでの使用例

・atcoder <a href="https://atcoder.jp/contests/abc140/tasks/abc140_f">ABC140F</a> <a href="https://atcoder.jp/contests/abc140/submissions/25074549">提出結果</a>
・atcoder <a href="https://atcoder.jp/contests/agc005/tasks/agc005_b">AGC005B</a> <a href="https://atcoder.jp/contests/agc005/submissions/25074519">提出結果</a>
・defaultdictとしての例: atcoder <a href="https://atcoder.jp/contests/diverta2019-2/tasks/diverta2019_2_b">diverta 2019 Programming Contest 2</a> <a href="https://atcoder.jp/contests/diverta2019-2/submissions/25074471">提出結果</a>
・atcoder <a href="https://atcoder.jp/contests/abc217/tasks/abc217_d">ABC217D</a> <a href="https://atcoder.jp/contests/abc217/submissions/25653023">提出結果</a>
・atcoder <a href="https://atcoder.jp/contests/abc217/tasks/abc217_e">ABC217E</a> <a href="https://atcoder.jp/contests/abc217/submissions/25653295">提出結果</a>
# 最後に
LGTMしていただけると大変励みになりますので参考になった方いましたらよろしくお願いいたします。