Friday, February 25, 2011

Prefix sums & Fenwick trees

Today I am going to post about prefix sums & an associated topic Fenwick trees ( aka Binary Indexed Trees or BIT in short ). This is one data structure I absolutely love and just for information BIT is so powerful that I use it in almost every third problem.

In this whole post I'll use Q for the number of queries to be answered  & N to denote input size.

Problem : We are given an array of numbers, say A.  A query consists of two indices i & j we need the sum A[i] + A[i+1]...A[j] quickly. An update consists of two numbers i & val such that you want A[i] to be incremented by val.

Solution 1 : Store a 2D table - Sums. Sums[i][j] stores the sum of A[i] + .. A[j].
This takes O(N^2) space, O(N^2) pre-process time and O(1) per query. Space requirement usually makes it virtually useless. Also we would shortly consider dynamic arrays as well.  Then for every update made to array, a lot of entries (O(N^2) infact ) are to be recomputed which makes it very slow.



Observation : If we can answer a simpler version of this problem where i = 1, we an answer the general version as well. A[1] + A[2] ... A[i] is called ith prefix sum of A, which I denote with Prefix[i].  Our query is just Prefix[j] - Prefix[i-1].


Solution 2 : So create a new array for prefix sums. Compute prefix sums for all i in O(N) time total using Prefix[i] = Prefix[i-1] + A[i].And now we can answer each query in constant time . But again if array is made dynamic, for every update we may have to change O(N) prefix sums. Too bad !


Solution 3 : Partition the array A in sqrt(N) segments each of size ( N / sqrt(N) ) = sqrt(N) and store sum for each segment. Now when you query an A[i] + ... + A[j] , first break it into prefix sum queries. Say now you want to answer Prefix[k]. So keep moving over segments which are completely in [1...k] adding their sums. There is utmost 1 incomplete segment. Move over each element of this which in included in range [1..k] . There are O(sqrt N) full segments you moved over and also O(sqrt N) elements you moved over in last incomplete segment. So query time is O(sqrt N). Also update time is O(sqrt N) as we have to update 1 segment only.
Easy to code, fast enough usually.

Solution 4 BIT. Note in prefix sums we store sum of a segment at each point in prefix array. That gives a good query time ( O(1) ) but large number of segments contains a given point which makes them bad for updates. As we saw if we make segments of smaller sizes such that each point in involved in only 1 segment, we can speed up the updates on cost of queries.  We need something in middle. Every element should be in more then 1 segment but not in too many either. Atleast  we do get a feel that a good solution would have to partition input array in segments and store sum for each.  So question to ask is what size (or form) segments ?

Solution to this is instead of using a fixed size segments , use a variable scheme . Suppose you have array A given. I will construct a new array called BIT which holds some segment sums.
In particular, if i has k trailing zeroes in binary representation,
BIT[i] = A[i] + A[i -1]... A[i - 2^k + 1].

i.e it stores sum of last 2^k entries of A. This looks random for sure. But its powerful. Lets see how.

i ) Querying using BIT : As before break query in prefix queries. Say we want to answer prefix query [1...i]. Now look at BIT[i]. It contains sum of a segment of A in particular of length 2^k where there are k trailing zeroes in binary representation of i. We can add these 2^k entries in O(1) time : BIT[i] is this sum by definition. What about other entries ? Well that sum is nothing but prefix query i - 2^k which we can find recursively. Question to ask is how fast is this ? Well it takes no more then O(log N) time. How ? Every recursive calls removes the right most 1 from binary representation of i. There are atmost log(N) digits in binary of N. Hence the time.

2) Updating BIT : Almost same . If you change A[i], then what all BIT values change ? BIT[i] changes for sure. What is the next j such that segment of BIT[j] contains i ? It would be one such element st it has k trailing zeroes and j - i <= 2^k.  Also we want smallest such j so that we can iterate over all in some order.
Convince yourself that this j is nothing but i + 2^k where k is number of trailing zeroes in i. Once you understand this, you are done as now when you change BIT[j] you propagate the update further exactly like this.

Now only 2 more things and you know it all :
1) Convince yourself that you actually dont need to store A anymore. Storing BIT array is sufficient for all purposes.
2) How to find 2^k where k is number of trailing zeros ? Again convince yourself that this nothing but : i & (-i)

Now lets throw in some code here :

// Gives prefix sum of segment [1...i]

int query(int i)   
{
 int ans = 0;
 while( i > 0)
 {
   ans += BIT[i];
   i -=  i & (-i);
 }
 return ans;
}


// Add val to A[i].

void update(int i, int val)
{
 while(i <= N)
 {
   BIT[i] += val;
   i += i & (-i);
 }
}




You see ? Under 10 lines of code you have both O(log N) updates and O(log N) queries. Power of this data structure is huge  and only more thought and practice can help you understand that.  This in practice is lightening fast with such a low constant involved in that O(log N).


I should link up some problems soon. Till then happy Binary indexing :)