【Python】バケット法（平方分割）

機能

セグ木と同様に、一点更新と区間取得ができます。ただし計算量はセグ木がの区間取得がO(logN)のところがバケット法だとO(\sqrt{N})かかり、遅くなります。基本的にセグ木と同じことしかできないにも関わらず計算量の点では性能は一段階劣るという感じなので出番はあまりないですが、セグ木に比べて構造が単純な分、問題に合わせて中身をいじったり、デバッグしたりはしやすいのではないかと考えています（今まで一度もそうした機会はないですが）。
また、演算の順番を崩さないので非可換モノイドにも対応可です。

使い方

• インスタンスの作成：基本的にはセグ木と同じで、配列の初期値（あれば）と二項演算と単位元を引数として渡します。逆元がある演算の場合は逆元をinvとして指定することで少し計算が早くなります。引数sizeを使ってバケットサイズを指定することもできます（バケットサイズを指定しない場合、Nの平方根に近い値がデフォルトで割り当てられます）
  • bk = Bucket(N,data=A,f=lambda a,b:a+b,e=0)
  • (逆元がある場合）bk = Bucket(N,data=A,f=lambda a,b:a+b,e=0, inv=lambda a,b:a-b)
  • (バケットサイズを指定する場合）bk = Bucket(N,data=A,f=lambda a,b:a+b,e=0, inv=lambda a,b:a-b, size=100)
• 一点更新（iは0-indexed)
  • bk.set_val(i, x)
• 区間取得（0-indexedかつ半開区間。range(l,r)と同じ）
  • bk.prod(l,r)
• 全区間取得：始点から終点までの全区間の計算結果が欲しいときはbk.prod(0,N)とするよりこちらの方が速い
  • bk.prod_all()

使用例

-セグ木の1.5倍程度の時間でAC
- https://atcoder.jp/contests/practice2/submissions/22311944
- 逆元が定義されている場合は更新作業が一回の演算で済むので早くなる。
- (+の逆元は-) https://atcoder.jp/contests/practice2/submissions/22312223
- (xorの逆元もxor) https://atcoder.jp/contests/abc185/submissions/22312514
- 非可換モノイドにも対応（タコヤキオイシクナール）
- https://atcoder.jp/contests/arc008/submissions/22313073
- ライブラリチェッカーでは一部TLE
- https://judge.yosupo.jp/submission/61631

参考にした記事など

プログラミングコンテストでのデータ構造 from Takuya Akiba

www.slideshare.net

最初に自力実装したデータ構造なので思い入れがあります。

class Bucket:
    def __init__(self, N, data=None, f=max, e=0, inv=None, size=None):
        self.N = N
        if not data: self.data = [e for _ in range(N)]
        else: self.data = data
        if not size: size = int(N**0.5)
        self.size, self.n = size,(N+size-1)//size
        self.f,  self.e, self.inv = f, e, inv
        self.middle = [e for _ in range(self.n)]
        if data:
            for i in range(self.n):
                for j in range(i*size, min((i+1)*size, N)): self.middle[i] = self.f(self.middle[i], self.data[j])
        self.all = e
        if data:
            for i in range(self.n): self.all = self.f(self.all, self.middle[i])
                
    def set_val(self, i, x):
        # 逆元が定義されている場合は一回の演算で済ませる。
        if self.inv:
            self.middle[i//self.size] = self.inv(self.middle[i//self.size], self.data[i])
            self.middle[i//self.size] = self.f(self.middle[i//self.size], x)
            self.all = self.inv(self.all, self.data[i])
            self.all = self.f(self.all, x)
            self.data[i] = x
            return
        self.data[i] = x
        ib = i//self.size
        self.middle[ib] = self.e
        for j in range(ib*self.size, min((ib+1)*self.size, N)): self.middle[ib] = self.f(self.middle[ib], self.data[j])
        self.all = self.e
        for i in range(self.n): self.all = self.f(self.all, self.middle[i])
            
    # [L, R)
    def prod(self, l, r):
        ibl, ibr = l//self.size, (r-1)//self.size
        out = self.e
        if ibl==ibr:
            for i in range(l, r): out = self.f(out, self.data[i])
            return out
        for i in range(l, (ibl+1)*self.size): out = self.f(out, self.data[i])
        for i in range(ibl+1, ibr): out = self.f(out, self.middle[i])
        for i in range(ibr*self.size, r): out = self.f(out, self.data[i])
        return out

    def prod_all(self): return self.all