Study notes: Extensions of binary indexed trees

Exploring the potentials of binary indexed trees (BIT)

Originally posted my Luogu blog [in Simplified Chinese]. This is an translated version in English.

Background Knowledge Recap

A Binary Indexed Tree (BIT), also known as a Fenwick Tree, stores at each index i the sum of a segment of length lowbit(i). This allows for \(O(\log N)\) time complexity operations for both single-point updates and prefix sum queries. Below is a sample implementation (thanks to @sjkmost for providing the code):

#define lowbit(x) ((x) & -(x))

template <typename T, int N>
struct BIT {
    T c[N];
    int n;

    void modify(int i, const T &rhs) { // Single-point update
        for (; i <= n; i += lowbit(i))
            c[i] = c[i] + rhs;
    }

    T query(int i) { // Prefix sum query
        T k = T();
        for (; i; i -= lowbit(i))
            k = k + c[i];
        return k;
    }
};

Extension #1: Range Update, Point Query

The basic functionality of a BIT is single-point update and prefix sum query. But how can we adapt it to support range updates and point queries?

It’s not too complicated—this can be done using a difference array dif. In a difference array, the value at any single point is equivalent to the prefix sum at that point. So, to update a range \([l..r]\), you simply update l by +c and r + 1 by -c.

Here’s how the implementation looks:

if (opt == 1) { // Range update
    a = read(), b = read(), c = read();
    dif.modify(a, c);
    dif.modify(b + 1, -c);
} else { // Point query
    pint(dif.query(read()));
    pchar('\n');
}

Extension #2: Range Update, Range Query

At first glance, implementing range updates and range queries using a BIT might seem impossible. I struggled with this for a long time—until @sjkmost gave me a brilliant insight.

Let’s break it down.

Suppose we use a BIT to store a difference array dif. What properties does the prefix sum at index ( i ) have?

\[\displaylines{ a_1 = & \text{dif}_1 \\\ a_2 = & \text{dif}_1 + \text{dif}_2 \\\ a_3 = & \text{dif}_1 + \text{dif}_2 + \text{dif}_3 \\\ \vdots \\\ a_i = & \sum_{j=1}^{i} \text{dif}_j }\]

If we observe carefully, when computing the sum of values from 1 to \(i\), index 1 is added \(i\) times, index 2 is added \(i - 1\) times, index 3 is added \(i - 2\) times, …, and index \(i\) is added once. It forms a pyramid shape of coefficients.

But this pattern is hard to maintain directly.

Now here’s the trick: If we multiply each \(a_i\) by \(i\), we get:

\[ia_i = \sum_{j=1}^{i} i \cdot \text{dif}_j\]

Now reverse the logic: If you maintain both a regular difference array (dif1) and another one to track \(i \cdot \text{dif}_i\) (dif2), you can reconstruct both the total value and its weighted sum.

Then, to get the sum over a range \([x..y]\), we compute:

\[\sum_{j=x}^{y} a_j = \left( \text{dif1.query}(y) \cdot (y+1) - \text{dif2.query}(y) \right) - \left( \text{dif1.query}(x-1) \cdot x - \text{dif2.query}(x-1) \right)\]

Code

if (opt == 1) { // Range update
    k = read();
    dif1.modify(x, k);
    dif1.modify(y + 1, -k);
    dif2.modify(x, k * x);
    dif2.modify(y + 1, -k * (y + 1));
} else { // Range query
    pint(
        (dif1.query(y) * (y + 1) - dif2.query(y)) -
        (dif1.query(x - 1) * x - dif2.query(x - 1))
    );
    pchar('\n');
}

Summary

The Binary Indexed Tree (BIT) is a very practical data structure, known for its small constant factors, low code complexity, and its robustness against common errors. Although its basic functionality might seem somewhat limited, with enough thought and innovation, many extended applications can be developed. To excel in competitive programming, one must think critically and innovate continuously.