用樹狀數組解決區間查詢問題

       本文擴寫自郭神的《樹狀數組新應用》，在此表示膜拜。樹狀數組的學名貌似叫做Binary Index Tree，關於它的基本應用可參考Topcoder上的這篇Tutorial.

樹狀數組可以看作一個受限制的線段樹，它維護一個數組，最經典的樹狀數組支援的基本操作有兩個：(1)改變某一個元素的值 (2)查詢某一個區間內所有元素的和。在此基礎上，經過簡單的變形可以變成支援另一組操作：(1)把一個區間內所有元素都加上一個值 (2)查詢某一個元素的值。這兩個都是已經泛濫了的東西了，在此不贅述。

簡單的樹狀數組模型是不支援這樣一組操作的：(1)把某一個區間內所有元素都加上一個值 (2)查詢某一個區間內所有元素的和。當然，這個東西可以用線段樹完成，但是線段樹占記憶體比較大，寫起來也比較繁（對我這種不會資料結構的人而言）。下面我們用一個改進版的樹狀數組完成這個任務。

首先一個觀察是區間操作總可以變成從最左端開始，比如把區間[3..6]都加10，可以變成[1..6]加10, [1..2]減10。查詢也類似。於是下面只關心從最左端開始的情況。定義Insert(p, d)表示把區間[1..p]都加d，Query(p)表示查詢區間[1..p]之和。

我們考慮調用一次Insert(p, d)對以後的某次查詢Query(q)的影響：

(1) 如果p<=q，總的結果會加上p*d (2) 如果p>q，總的結果會加上q*d

也就是說，Query(q)的結果來源可分為兩部分，一部分是Insert(p1,d) (p1<=q)，一部分是Insert(p2,d) (p2 > q)。我們用兩個數組B[], C[]分別維護這兩部分資訊，B[i]表示區間右端點恰好是i的所有區間的影響之和，C[i]表示區間右端點大於i的所有區間的影響之和。每當遇到 Insert時，考慮當前的Insert會對以後的Query產生什麼影響，更新B和C數組；當遇到Query時，把兩部分的結果累加起來。

具體來說，當我們遇到Insert(p, d)時，把B[p]增加p*d，把C[1], C[2], …, C[p-1]都增加d。當遇到Query(p)時，查詢B[1]+B[2]+…+B[p]+C[p]*p即可。可以發現對B數組是修改單個元素，查詢區間和；對C數組是修改區間，查詢單個元素，這恰好對應於一開始說的樹狀數組支援的基本操作。於是我們用兩個樹狀數組漂亮地完成了任務。

範例程式碼：

#include <cstdio>const int MAXN = 1024;int B[MAXN], C[MAXN];#define LOWBIT(x) ((x)&(-(x)))void bit_update(int *a, int p, int d){for ( ; p && p < MAXN ; p += LOWBIT(p))a[p] += d;}int bit_query(int *a, int p){int s = 0;for ( ; p ; p -= LOWBIT(p))s += a[p];return s;}void bit_update2(int *a, int p, int d){for ( ; p ; p -= LOWBIT(p))a[p] += d;}int bit_query2(int *a, int p){int s = 0;for ( ; p && p < MAXN ; p += LOWBIT(p))s += a[p];return s;}inline void _insert(int p, int d){bit_update(B, p, p*d);bit_update2(C, p-1, d);}inline int _query(int p){return bit_query(B, p) + bit_query2(C, p) * p;}inline void insert_seg(int a, int b, int d){_insert(a-1, -d);_insert(b, d);}inline int query_seg(int a, int b){return _query(b) - _query(a-1);}int main(){int com, a, b, c;while (scanf("%d%d%d",&com,&a,&b) != EOF)    {a += 2; b += 2;//防止出現負數if (com == 0){//更新scanf("%d",&c);insert_seg(a, b, c);}        else        {//查詢printf("%d\n",query_seg(a,b));}}return 0;}