Trees

Fenwick Tree / Binary Indexed Tree (BIT)

NOTE used in log-linear attention :O

A Fenwick Tree stores partial sums of an array so that prefix sums and point updates are both fast.

It supports:

\[ \texttt{prefix}(r)=\sum_{i=1}^{r} A[i] \]

and

\[ A[i] \leftarrow A[i]+\Delta \]

in $O(\log n)$ time, using $O(n)$ space.

Use 1-indexing. Define $\operatorname{lowbit}(i)=i\&(-i)$, the largest power of two dividing \(i\).

Each BIT entry stores

\[ F[i]=\sum_{k=i-\operatorname{lowbit}(i)+1}^{i} A[k]. \]

or \(F[i]\) stores the sum of the block ending at \(i\) with length \(\operatorname{lowbit}(i)\).

Prefix query moves downward:

\[ i \leftarrow i-\operatorname{lowbit}(i) \]

Point update moves upward:

\[ i \leftarrow i+\operatorname{lowbit}(i) \]

Range sums come from prefix subtraction:

\[ \sum_{k=l}^{r} A[k]= \texttt{prefix}(r)-\texttt{prefix}(l-1). \] \[ \text{query},\text{ update},\text{ range sum}: O(\log n) \]

CODE

Class BIT:


def init(self, n):
    self.n = n
    self.bit = [0] * (n + 1)
    
def update(self, i, delta):
    while i <= self.n:
        self.bit[i] += delta
        i += i & -i
    
def prefix(self, i):
    s = 0
    while i > 0:
        s += self.bit[i]
        i -= i & -i
    return s
    
def range_sum(self, l, r):
    return self.prefix(r) - self.prefix(l - 1)