Monday, May 27, 2013

Fenwick tree range updates

A Fenwick tree is a wonderful data structure that supports two operations on an array: increment a given value by a given amount, and find the sum of a segment of values, both in O(log n) time. What's more, the Fenwick tree is represented as just an array of the same size as the array being updated and queried, and we don't need to store the original array itself! In other words, we can support the above operations without any additional memory. Yesterday I've discovered that it's capable of even more!

For a start, here is the entire code for a Fenwick tree:

private void update(int at, int by) {
    while (at < data.length) {
        data[at] += by;
        at |= (at + 1);
    }
}

private int query(int at) {
    int res = 0;
    while (at >= 0) {
        res += data[at];
        at = (at & (at + 1)) - 1;
    }
    return res;
}

Here, query(at) returns the sum of all elements from the first element to at-th element. Finding a sum of arbitrary segment can be done via query(right)-query(left-1). You can find more details using your favorite search engine.

The standard Fenwick tree only supports updating one element at a time. However, it is quite natural to expect it to handle range updates, too - the update command now becomes "increment each value between left-th and right-th by the given amount". It turns out that we can modify the tree to handle such updates in O(log n) time, too! Here's how (warning, code untested - but I hope it works):

private void update(int left, int right, int by) {
    internalUpdate(left, by, -by * (left - 1));
    internalUpdate(right, -by, by * right);
}

private void internalUpdate(int at, int mul, int add) {
    while (at < dataMul.length) {
        dataMul[at] += mul;
        dataAdd[at] += add;
        at |= (at + 1);
    }
}

private int query(int at) {
    int mul = 0;
    int add = 0;
    int start = at;
    while (at >= 0) {
        mul += dataMul[at];
        add += dataAdd[at];
        at = (at & (at + 1)) - 1;
    }
    return mul * start + add;
}

In other words, we implement a Fenwick tree with range updates via a normal (point-update) Fenwick tree that stores linear functions instead of just values. Is this a well-known trick?

2 comments: