一般树状数组

单点更改, 区间查询

class FenwickTree:
    def __init__(self, size):
        self.size = size
        self.tree = [0] * (self.size + 1)

    def update(self, index, delta):
        while index <= self.size:
            self.tree[index] += delta
            index += index & -index

    def query_prefix(self, index):
        res = 0
        while index > 0:
            res += self.tree[index]
            index -= index & -index
        return res

    def query(self, index):
        return self.query_range(index, index)

    def query_range(self, left, right):
        return self.query_prefix(right) - self.query_prefix(left - 1)

例题

P3374 【模板】树状数组 1

差分树状数组

区间更改, 单点查询

class DiffFenwickTree:
    def __init__(self, size):
        self.size = size
        self.tree = [0] * (self.size + 1)

    def _update(self, index, delta):
        while index <= self.size:
            self.tree[index] += delta
            index += index & -index

    def update(self, index, delta):
        self._update(index, delta)
        self._update(index + 1, -delta)

    def update_range(self, start, end, delta):
        self._update(start, delta)
        self._update(end + 1, -delta)

    def query(self, index):
        res = 0
        while index > 0:
            res += self.tree[index]
            index -= index & -index
        return res

例题

P3368 【模板】树状数组 2

双树状数组

水群知道的, 比较有趣的知识.

支持区间修改区间查询

class RURQFenwickTree:
    def __init__(self, size):
        self.ft1 = FenwickTree(size)
        self.ft2 = FenwickTree(size)

    def _update(self, index, delta):
        self.ft1.update(index, delta)
        self.ft2.update(index, delta * index)

    def update(self, index, delta):
        self.update_range(index, index, delta)

    def update_range(self, left, right, delta):
        self._update(left, delta)
        self._update(right + 1, -delta)

    def query_prefix(self, index):
        return (index + 1) * self.ft1.query_prefix(index) - self.ft2.query_prefix(index)

    def query_range(self, left, right):
        return self.query_prefix(right) - self.query_prefix(left - 1)

    def query(self, index):
        return self.query_range(index, index)

例题

P3372 【模板】线段树 1 题解提交

证明

差分数组 $b_i$ 和原数组 $a_i$ 的关系

在代码中, $a_i$ 是原数组, 而 $b_i$ 是它的差分数组, 定义为：

$b_i = a_i - a_{i-1}$

这样，原数组 $a_i$ 可以表示为差分数组 $b_i$ 的前缀和：

$a_i = \sum_{j=1}^{i} b_j$

区间和 `sum(a[l:r])` 的推导

现在考虑计算区间 [l, r] 的和：

$\sum_{i=l}^{r} a_i = \sum_{i=l}^{r} \sum_{j=1}^{i} b_j$

这个双重求和可以重新排列：

$\begin{aligned} \sum_{i=l}^{r} \sum_{j=1}^{i} b_j &= \sum_{j=1}^{r} b_j \cdot (r - j + 1) - \sum_{j=1}^{l-1} b_j \cdot (l - j) \end{aligned}$

对于以上操作, 更直观的理解, 考虑由 $b_j$ 组成的”三角形”, 由上倒下分别:

$\begin{array}{c} &b_1 \cr &b_1, b_2 \cr &\vdots \cr &b_1, b_2, \cdots, b_n \end{array}$

可以发现, 期间有 $n$ 个 $b_1$ , $n-1$ 个 $b_2$ 直到一个 $b_n$ . 所以式子可以如此转化:

$\sum_{i=1}^{n} \sum_{j=1}^{i} b_j = \sum_{j=1}^{n} b_j \cdot (n - j + 1)$

当只需要 $l \to r$ 时, 相当于大三角形减去小三角形:

$\begin{array}{c} &b_1, \cdots, b_l \cr &b_1, b_2, \cdots, b_{l+1} \cr &\vdots \cr &b_1, b_2, b_3, \cdots, b_n \end{array}$

可以得到:

$\sum_{i=l}^{r} \sum_{j=1}^{i} b_j = \sum_{i=1}^{r} \sum_{j=1}^{i} b_j - \sum_{i=1}^{l-1} \sum_{j=1}^{i} b_j = \sum_{j=1}^{r} b_j \cdot (r - j + 1) - \sum_{j=1}^{l-1} b_j \cdot (l - j)$

使用树状数组优化计算

为了高效计算这个和, 代码使用了两个树状数组:

tr1[i]: 维护差分数组 $b_i$ 的前缀和(即 sum(b[1:i]))
tr2[i]: 维护 i * b[i] 的前缀和(即 sum(j * b[j] for j in range(1,i+1)) 这样, 区间和可以表示为:

$\sum_{i=l}^{r} a_i = \text{query}(r) - \text{query}(l-1)$

其中 query(p) 计算的是：

$\text{query}(p) = (p + 1) \cdot \text{sum}(b[1:p]) - \text{sum}(j \cdot b[j] \text{ for j in range(1, p+1)})$

即:

$\text{query}(p) = (p + 1) \cdot tr1[p] - tr2[p]$

到此, 我们完成了区间查询的证明.

Python 树状数组