Python版 AtCoder Library (Fenwick Tree)

この記事では、AtCoder Library (ACL)のfenwicktreeをPythonで書き直したものを公開しています。

Fenwick Tree

長さ の配列に対し、

• 要素の１点変更
• 区間の要素の総和

を で求めることが出来るデータ構造です。

コード

class fenwick_tree:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)
    
    def add(self, p, x):
        p += 1
        while p <= self.n:
            self.bit[p] += x
            p += p & -p
    
    def sum_(self, p):
        s = 0
        while p > 0:
            s += bit[p]
            p -= p & -p
        return s
    
    def sum(self, l, r):
        return sum_(r) - sum_(l)

コンストラクタ

fw = fenwick_tree(n)

• 長さの配列を作ります。初期値はすべてです。

計算量

• 

add

fw.add(p, x)

a[p] += x を行う

制約

• 

計算量

• 

sum

fw.sum(l, r)

a[l] + a[l + 1] + ... + a[r - 1] を返す。

制約

• 

計算量

• 

使用例

AC code of https://atcoder.jp/contests/practice2/tasks/practice2_b

############################################

class fenwick_tree:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)
    
    def add(self, p, x):
        p += 1
        while p <= self.n:
            self.bit[p] += x
            p += p & -p
    
    def sum_(self, p):
        s = 0
        while p > 0:
            s += self.bit[p]
            p -= p & -p
        return s
    
    def sum(self, l, r):
        return self.sum_(r) - self.sum_(l)

############################################

N, Q = map(int,input().split())
a = list(map(int,input().split()))

fw = fenwick_tree(N)

for i in range(N):
    fw.add(i, a[i])

for _ in range(Q):
    q = list(map(int,input().split()))
    if q[0] == 0:
        fw.add(q[1], q[2])
    else:
        print(fw.sum(q[1], q[2]))

参考文献

• 『プログラミングコンテストチャレンジブック [第2版]』
• https://atcoder.github.io/ac-library/document_ja/fenwicktree.html