树状数组

考虑一个数组，我们有两个需求：①区间更新；②区间和查询。

方案1

最朴素的想法，时间复杂度都是O(n)。

arr = [2,0,1,7,0,9]
for i in range(l, r):	# 区间更新
    arr[i] += value
for i in range(l, r):	# 区间和查询
    sum += arr[i]

方案2

高级点的想法，时间复杂度都是O(1)。这里用到了前缀数组和差分数组。

arr = [2,0,1,7,0,9]
prefix, diff = [2,2,3,10,10,19], [2,-2,1,6,-7,9]
diff[l] += value	# 区间更新
diff[r] -= value
sum = prefix[r] - prefix[l]	# 区间和查询

还有点小问题，区间和查询的O(1)是建立在arr不会更新的前提上，如果arr是动态更新的，那么arr每更新一次，就需要重建prefix，代价是O(n)的时间复杂度。

树状数组

下面代码给出了树状数组的初始化，arr和tree的下标从1开始。

arr  = [0,2,0,1,7,0,9]
tree = [0,0,0,0,0,0,0]

def add(i, value):
    while i < len(arr):
        tree[i] += value
        i += i&(-i)

for i in range(1, len(arr)):
    add(i, a[i])

print tree	# [0,2,2,1,10,0,9]

i=i&(-i)将保留i的最低位的1，其余所有位置0。

那么我们如何基于树状数组做区间和查询？

def ask(i):
    while i:
        sum += tree[i]
        i -= i&(-i)
    return sum

for i in range(1, len(arr)):
    print ask(i)	# 2 2 3 10 10 19

sum = ask(r) - ask(l)

add()时间复杂度为O(log(n))，ask()时间复杂度为O(log(n))。

若以add()做区间更新，承担O(nlog(n))的时间复杂度。

for i in range(l, r):
    add(i, value)

上面提到了区间更新可以使用差分数组优化时间，但是树状数组却是基于原数组建立的，如果我们基于原数组的差分数组建立树状数组，会是什么样子？

arr = [0,2,0,1,7,0,9]
diff, tree = [0,2,-2,1,6,-7,9], [0,0,0,0,0,0,0]

for i in range(1, len(arr)):
    add(i, diff[i])

print tree	# [0,2,0,1,7,-7,2]

for i in range(1, len(arr)):
    print ask(i)	# 2 0 1 7 0 9

可以看到基于差分数组建立树状数组，ask(i)的返回值为arr[i]，因此该方法适用于区间更新，单点查询，而上面的方法适用于单点更新，区间查询。