オフラインクエリのset実装

この記事はpythonで競プロやってる~~ひねくれた~~人向けです。

最近ボンタン飴がマイブームなまつかわです、
この前のCPSCOの記事
CPSCO2019参加記(1日目) - Tallfallの日記

でpythonにstd::setをくれと書いたらなんとそれっぽいのをnadareさんとfineさんが実装してくれたので記事にまとめてみることにします。
というよりは自分のための実装のメモに近いです。
参照: Submission #5671536 - CPSCO2019 Session1

やりたいことは主に以下の３つ

add(i, x) : iを集合にx個挿入(削除)する
get(v) : 集合の要素中v番目に小さい値を取得
count(i) : 集合の要素であってi以下(未満)の個数を取得(iが何番目か分かる)

これに加えてv番目までの和とかi以下の和とかを求めるクエリなどにも応用できる

これらを扱いたい値全体の要素数を $n$ として $O(\log{n})$ で処理する

実装にはBIT上の二分探索を使用している
分かりやすい記事をけんちょんさんが書いているので未読の方は読んだ方が以下が分かりやすいかと思います。
k 番目の値を高速に取り出せるデータ構造のまとめ - BIT上二分探索や平衡二分探索木など - Qiita

以下メモ

クラス名は参照元に従ってDammyMapとしてます
挿入もしくは削除する値は予め分かっている必要があります
分かりやすいようにaddで座圧したりgetで戻したりして値自体で扱えるようにしてますが、どうしてもTLEしそうとかだったらindexでやり取りするのが速いと思います

from itertools import accumulate
from collections import Counter
from bisect import bisect as br, bisect_left as bl
class PMS:
    #1-indexed
    def __init__(self, A, B, issum = False):
        #Aに初期状態の要素をすべて入れる,Bは値域のリスト
        self.X, self.comp = self.compress(B)
        self.size = len(self.X)
        self.tree = [0] * (self.size + 1)
        self.p = 2**(self.size.bit_length() - 1)
        self.dep = self.size.bit_length()
        
        CA = Counter(A)
        S = [0] + list(accumulate([CA[self.X[i]] for i in range(self.size)]))
        for i in range(1, 1+self.size):
            self.tree[i] = S[i] - S[i - (i&-i)]
        if issum:
            self.sumtree = [0] * (self.size + 1)
            Ssum = [0] + list(accumulate([CA[self.X[i]]*self.X[i] for i in range(self.size)]))
            for i in range(1, 1+self.size):
                self.sumtree[i] = Ssum[i] - Ssum[i - (i&-i)]
    
    def compress(self, L):
        #座圧
        L2 = list(set(L))
        L2.sort()
        C = {v : k for k, v in enumerate(L2, 1)}
        # 1-indexed
        return L2, C
    
    def leng(self):
        #今入っている個数を取得
        return self.count(self.X[-1])
    
    def count(self, v):
        #v(Bの元)以下の個数を取得
        i = self.comp[v]
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s
    
    def less(self, v):
        #v(Bの元である必要はない)未満の個数を取得
        i = bl(self.X, v)
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s
    
    def leq(self, v):
        #v(Bの元である必要はない)以下の個数を取得
        i = br(self.X, v)
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s
 
    def add(self, v, x):
        #vをx個入れる,負のxで取り出す,iの個数以上取り出すとエラーを出さずにバグる
        i = self.comp[v]
        while i <= self.size:
            self.tree[i] += x
            i += i & -i

    def get(self, i):
        # i番目の値を取得
        if i <= 0:
            return -1
        s = 0
        k = self.p
        for _ in range(self.dep):
            if s + k <= self.size and self.tree[s+k] < i:
                s += k
                i -= self.tree[s]
            k //= 2
        return self.X[s]
    
    def addsum(self, i, x):
        #sumを扱いたいときにaddの代わりに使う
        self.add(i, x)
        x *= i
        i = self.comp[i]
        while i <= self.size:
            self.sumtree[i] += x
            i += i & -i
    
    def countsum(self, v):
        #v(Bの元)以下のsumを取得
        i = self.comp[v]
        s = 0
        while i > 0:
            s += self.sumtree[i]
            i -= i & -i
        return s
    
    def getsum(self, i):
        #i番目までのsumを取得
        x = self.get(i)
        return self.countsum(x) - x*(self.count(x) - i)