一文搞懂树状数组

发布于 2022-02-17 # 算法与数据结构

介绍

树状数组是由Peter M. Fenwick提出的二叉索引树（Binary Indexed Trees）结构，最初用于数据压缩。在算法竞赛中，常用于区间操作。

类似的数据结构还有线段树，线段树可以实现树状数组所有的操作，甚至更多。而树状数组代码简洁，运行速度也线段树快，占用内存空间也比线段树少，如果是一个单点修改的问题，树状数组绝对是一个不二选择。

接下来我们引入一个问题：

已知有 $n$ 个箱子，你需要进行以下两种操作：

将 $k$ 个大理石放入第 $i$ 箱子中

求第 $l$ 到 $r$ 个箱子中大理石总共有多少个

假设我们要进行 $m$ 次操作。

对于这个问题，很容易想到朴素的做法：直接修改数组、遍历求和。

朴素做法对于操作 1 的时间复杂度为： $O(1)$ 。操作 2 时间复杂度为： $O(n)$ 。当所有操作都为 2 的时候，达到最坏时间复杂度 $O(nm)$ 。

如果数据量大的话，使用朴素的做法显然会 $TLE$ 。

这是道简单的单点修改 & 区间查询的题目，我们使用树状数组，查询的时间复杂度就可以降到 $O(logn)$ ，最坏时间复杂度就可以降到 $O(m logn)$ 了。

原理

下图为树状数组的原理：

是不是看起来非常抽象？完全无法理解其中是什么含义？是的就是非常的抽象，让我来解释一下这张图你就明白了~

树状数组的结构和线段树很类似，都是用一个大的节点去管理若干个小节点，查询的时候只需要查大节点就可以得到区间的信息。

图中 8 个蓝色的方块代表着数组 $a$ ，而绿色方块代表数组 $t$ ，即管理着数组 $a$ 的逻辑结构，为什么叫逻辑结构呢，我的理解是这个结构是根据逻辑虚构出来的，在实际内存中并不存在数组 $t$ ，而图中的 $t$ 则是将抽象的逻辑结构形象化了，有助于去理解数据结构。可以从图中看出：

$t[2]$ 管理的是 $a[1],a[2]$
$t[4]$ 管理的是 $a[1],a[2],a[3],a[4]$
$t[6]$ 管理的是 $a[5],a[6]$
$t[8]$ 管理的是 $a[1],a[2],...,a[7],a[8]$

以 t[6] 为例，方块的下半部分 $(0110)_2$ 是数组 索引的二进制 表示形式，而高亮部分的位（10）就是 lowbit。我们可以发现，lowbit 的值就是结点的覆盖长度，6的 $lowbit$ 值为 2，那么就可以得知t[6]结点的覆盖长度为 2。知道了覆盖长度，我们就可以利用覆盖长度来找上一个结点或下一个结点（父结点）了。

比如我们要找 t[6] 的父结点 t[8]，只需要t[6+lowbit(6)]就行了，由此我们可以得出: $t[i]$ 的父节点为 $t[i+lowbit(i)]$ 。

那如果要找 t[6] 的上一个结点 t[4] 呢？也很简单，和找父结点类似，索引值减去 lowbit 的值就可以了，也就是t[6-lowbit(6)]，由此我们可以得出: $t[i]$ 的上一个结点为 $t[i-lowbit(i)]$ 。

好了，知道了怎么找到上一个结点和父结点，那么我们就可以正式开始搞树状数组了~

等等，到这里你一定很懵逼，上面说的 $lowbit$ 到底是啥玩意？那怎么去求 $lowbit$ 呢？那么在正式开搞树状数组前，我们先详细的讲讲 $lowbit$ 运算~

lowbit运算

为了简洁起见，我们定义" $lowbit$ "为非负整数中最小的非零有效位，即非负整数 $n$ 在二进制表示下最低位 1 及其后面的 0 构成的数值。

例如，我们对十进制数字 44 进行 $lowbit$ 运算： $lowbit(44)=lowbit((101100)_2)=(100)_2=4$ 。

其中 $(100)_2$ 就是最低位 1 及其后面的 0 构成的数值。

但怎么通过程序去找到所谓的 $lowbit$ 数值呢？转成二进制再通过遍历寻找？这显然效率太低，不妨我们用位运算试试吧~

首先我们知道44的二进制为 101100，即 $(44)_{10}=(101100)_2$ 。
将 101100 按位取反，得到 010011，再将取反后的值 + 1，得到 010100。
观察101100 和 010100，可以看出，除了最低位的 1 和后面的 0，其余位上两者均不同。这时聪明的你可能已经发现了，将两者进行按位与运算不就可以得到lowbit了吗？没错，就是这样的~

简言之，就是将该二进制值和该二进制取反后+1 的值进行按位与运算，就可以得到 lowbit 值了，即：

lowbit(n)=n \& (\sim n+1)

众所周知，计算机中的有符号整数（比如int, long等）在底层是用补码表示的，而补码的规则是：

正数的补码就等于其原码
负数的补码等于其「原码的数值位逐位取反」然后「加一」

那么我们把正的 $n$ 改成 $-n$ ，不就等同于做「将 $n$ 取反后加一」这个操作吗？例如，补码形式的 44 为 $(0,101100)$ ，-44 的补码形式为 $(1,010011)$ ，则

(44)_{10} \& (-44)_{10} = (0,101100)_补\&(1,010011)_补=(0,000100)_补=(4)_{10}

所以最终可定义 $lowbit$ 运算为：

lowbit(n)=n\&-n

代码如下：

// Java
public int lowbit(int n) {
  	return n & -n;
}

单点修改 & 区间查询

单点修改

单点修改我们只需要将 $a_i$ 加上 $k$ , 更新 $a_i$ 所有的上级（父结点）即可。

例如对 $a_2 + 3$ ，具体过程如下图

代码如下：

// Java
public void add(int i, int k) {
  	while(i <= n) {
      	a[i] += k;
      	i += lowbit(i);
    }
}

区间查询

想要在单点修改中进行区间查询，首先要知道怎么求 $a$ 的前缀和。

还是利用 $lowbit$ 函数，一直找上一个结点进行求和。

例如求 $a_7$ 前缀和 $\sum_{i=1}^{7}$ ，过程如下图：

知道了怎么求前缀和后，求区间就很简单了，只需要将两个前缀和相减，即：

$\sum_{i=l}^{r}=\sum_{i=1}^{r}-\sum_{i=1}^{l-1}$

例如求区间和 $[3,7]$ ，过程如下图：

代码如下：

// Java
public int getPrefixSum(int i) {
  	int sum = 0;
  	while(i > 0) {
      	sum += a[i];
      	i -= lowbit(i);
    }
  	return sum;
}

public int getSum(int l, int r) {
  	return getPrefixSum(r) - getPrefixSum(l - 1);
}

区间修改 & 单点查询

区间修改

在前面的部分，我们已经讨论了如何通过树状数组实现单点修改和区间查询。接下来，我们将讨论如何通过树状数组实现区间修改和单点查询。

区间修改的核心思想是利用差分数组。假设我们有一个数组 $a$ ，我们可以构造一个差分数组 $d$ ，其中 $d[i] = a[i] - a[i-1]$ （ $d[1] = a[1]$ ）。通过差分数组，我们可以将区间修改操作转化为单点修改操作。

例如，如果我们想要将区间 $[l, r]$ 中的每个元素都加上 $k$ ，我们只需要在差分数组 $d$ 中进行如下操作：

将 $d[l]$ 加上 $k$
将 $d[r+1]$ 减去 $k$

这样，当我们查询某个位置 $i$ 的值时，只需要求差分数组 $d$ 的前缀和即可。

代码如下：

// Java
public void rangeAdd(int l, int r, int k) {
    add(l, k);
    add(r + 1, -k);
}

单点查询

单点查询的实现非常简单，只需要求差分数组 $d$ 的前缀和即可。由于我们已经通过树状数组维护了差分数组 $d$ ，因此单点查询的时间复杂度为 $O(\log n)$ 。

代码如下：

// Java
public int getValue(int i) {
    return getPrefixSum(i);
}

区间修改 & 区间查询

区间修改

在前面的部分，我们已经讨论了如何通过树状数组实现区间修改和单点查询。接下来，我们将讨论如何通过树状数组实现区间修改和区间查询。

为了实现区间修改和区间查询，我们需要维护两个树状数组：一个用于维护差分数组 $d$ ，另一个用于维护 $i \cdot d[i]$ 的前缀和。通过这两个树状数组，我们可以高效地实现区间修改和区间查询。

区间修改的实现与之前类似，我们只需要在差分数组 $d$ 中进行单点修改即可。

代码如下：

// Java
public void rangeAdd(int l, int r, int k) {
    add(l, k);
    add(r + 1, -k);
}

区间查询

区间查询的实现稍微复杂一些。我们需要利用两个树状数组来计算区间和。具体来说，区间 $[l, r]$ 的和可以通过以下公式计算：

\sum_{i=l}^{r} a[i] = \sum_{i=1}^{r} a[i] - \sum_{i=1}^{l-1} a[i]

而 $\sum_{i=1}^{x} a[i]$ 可以通过以下公式计算：

\sum_{i=1}^{x} a[i] = (x + 1) \cdot \sum_{i=1}^{x} d[i] - \sum_{i=1}^{x} i \cdot d[i]

因此，我们需要维护两个树状数组，一个用于维护 $\sum_{i=1}^{x} d[i]$ ，另一个用于维护 $\sum_{i=1}^{x} i \cdot d[i]$ 。

代码如下：

// Java
public int rangeSum(int l, int r) {
    return prefixSum(r) - prefixSum(l - 1);
}

private int prefixSum(int x) {
    return (x + 1) * getPrefixSum1(x) - getPrefixSum2(x);
}

private int getPrefixSum1(int x) {
    int sum = 0;
    while (x > 0) {
        sum += tree1[x];
        x -= lowbit(x);
    }
    return sum;
}

private int getPrefixSum2(int x) {
    int sum = 0;
    while (x > 0) {
        sum += tree2[x];
        x -= lowbit(x);
    }
    return sum;
}

BLOG

Josh's Blog

一文搞懂树状数组

介绍

原理

lowbit运算

单点修改 & 区间查询

单点修改

区间查询

区间修改 & 单点查询

区间修改

单点查询

区间修改 & 区间查询

区间修改

区间查询

练习题

参考