Fenwick treeで更新つき累積和

累積和は変更に弱い。Fenwick tree は区間を2の冪サイズに分け、必要な集計だけを更新します。

言語：Python / 計算量：更新・質問 O(log N)

考え方

内部は1始まり。i & -i で担当区間の長さを得ます。
加算更新は i += i & -i。
先頭和は i -= i & -i で集めます。

具体例

[2,5,1] の先頭2個は7。添字1に3を足すと[2,8,1]、先頭2個は10です。

実装

n=3;bit=[0]*(n+1)
def add(i,x):
    i+=1
    while i<=n:
        bit[i]+=x;i+=i&-i
def prefix(r):
    s=0
    while r>0:
        s+=bit[r];r-=r&-r
    return s
for i,x in enumerate([2,5,1]):add(i,x)
add(1,3)
print(prefix(2))

注意する条件

代入更新をしたいなら新値−旧値を加算します。内部添字0で更新すると無限ループです。

確認問題

[2,5,1] の添字1に3を加えた全体和は？

解答と理由

2+8+1=11です。

実装課題

講義の Fenwick 実装を使い [1,2,3,4] の添字2に5を加え、[1,4) の和を出す。入力なし。

出力:
14

参考実装

n=4;bit=[0]*(n+1)
def add(i,x):
    i+=1
    while i<=n:
        bit[i]+=x;i+=i&-i
def pref(r):
    s=0
    while r:
        s+=bit[r];r-=r&-r
    return s
for i,x in enumerate([1,2,3,4]):add(i,x)
add(2,5)
print(pref(4)-pref(1))

読了の記録・下書き・メモへ