树状数组详解

单点更新，区间查询

假设有一个数组，对他大量的修改和查询，修改的是数组中某一个元素的值，查询的是数组中任意一个区间的和。对于修改比较简单，时间复杂度是 O(1) ，而查询的时间复杂度是 O(n) 。有同学可能会说使用前缀和来优化，前缀和查询的时间复杂度确实是 O(1) ，但如果我们修改某一个元素的时候，前缀和后面的值也都要修改，时间复杂度是 O(n) 。那么我们综合一下，有没有一种方式可以让修改和查询时间复杂度降一个数量级呢？有的，那就是树状数组，他的修改和查询时间复杂度都是 O(logn) ，综合来看还是不错的。如下图所示，他就是一个树状数组，其中数组 a[] 是原始数组，数组 c[] 是树状数组。

alt

我们令树状数组每个位置存的是子节点值的和，则有：

c[1] = a[1];
c[2] = c[1]+a[2] = a[1]+a[2];
c[3] = a[3];
c[4] = c[2]+c[3]+a[4] = a[1]+a[2]+a[3]+a[4];
c[5] = a[5];
c[6] = c[5]+a[6] = a[5]+a[6];
c[7] = a[7];
c[8] = c[4]+c[6]+c[7]+a[8] = a[1]+a[2]+a[3]+a[4]+a[5]+a[6]+a[7]+a[8];
c[9] = a[9];
c[10] = c[9]+a[10] = a[9]+a[10];

通过上面的公式可以发现一个规律就是：

//  k 为 i 的二进制中最右边连续 0 的个数。
c[i] = a[i - 2^k+1] + a[i - 2^k+2] + ... + a[i];

比如 c[6] ， 6 的二进制是 (110) ，最右边有 1 个 0 ，那么 k 就等于 1 ，所以：

c[6] = a[6 - 2^1 + 1] + a[6] = a[5] + a[6]

在比如 c[4] ， 4 的二进制是 (100) ，最右边有 2 个连续的 0 ，那么 k 就等于 2 ，所以：

c[4] = a[4 - 2^2 + 1] + a[4 - 2^2 + 2] + a[4 - 2^2 + 3] + a[4] 
	 = a[1] + a[2] + a[3] + a[4]

通过上面的图我们可以发现，在树状数组 c[i] 中，如果 i 是奇数， c[i] 就在第 0 层，也就是树状数组的叶子节点，并且 c[i] = a[i] 。如果 i 不是奇数，那么在 i 的二进制位中，他右边有几个 0 ， c[i] 就在第几层。我们定义函数 int lowBit(int x) 表示只保留 x 的二进制中最右边的 1 ，其他位置全部变为 0 ，比如：

数字 12 的二进制是 00001100，则 lowBit(12) = 4（二进制是00000100）
数字 13 的二进制是 00001101，则 lowBit(13) = 1（二进制是00000001）
数字 14 的二进制是 00001110，则 lowBit(14) = 2（二进制是00000010）
数字 16 的二进制是 00010000，则 lowBit(16) = 16（二进制是00010000）

函数 int lowBit(int x) 的代码如下：

private int lowBit(int x) {
    return x & -x;
}

这个很好理解，比如数字 12 和 -12 ，他们的二进制如下，只要他们进行 & 运算就是我们想要的结果。

00000000 00000000 00000000 00001100 （12的二进制）
11111111 11111111 11111111 11110100 （-12的二进制）
00000000 00000000 00000000 00000100 （12和-12，&运算的结果）

令 s[i] 表示元素数组 a 的前 i 项和，通过上面的图可以找出 s 和 c 的关系如下：

s[1] = a[1] = c[1];
s[2] = a[1]+a[2] = c[2];
s[3] = a[1]+a[2]+a[3] = c[2]+c[3];
s[4] = a[1]+a[2]+a[3]+a[4] = c[4];
s[5] = a[1]+a[2]+a[3]+a[4]+a[5] = c[4]+c[5];
s[6] = a[1]+a[2]+a[3]+a[4]+a[5]+a[6] = c[4]+c[6];
s[7] = a[1]+a[2]+a[3]+a[4]+a[5]+a[6]+a[7] = c[4]+c[6]+c[7];
s[8] = a[1]+a[2]+a[3]+a[4]+a[5]+a[6]+a[7]+a[8] = c[8];

通过上面等式的关系我们可以得出：

s[i] = c[i] + c[i-k1] + c[i-k1-k2] + c[i-k1-k2-k3] + ……

这里的 k1，k2，k3，……，kn 都是 2 的 k 次方，实际上就是不断的抹去 i 的二进制中右边的 1 ，直到 i 变成 0 为止。比如数字 7 ，他的二进制是 111 ，所以 s[111]=c[111]+c[110]+c[100]（这里的数字是二进制表示），也就是 s[7]=c[7]+c[6]+c[4] 。

private int prefixSum(int i) {
    int sum = 0;
    while (i > 0) {
        sum += c[i];
        i -= lowBit(i);// 抹去二进制中最右边的 1 。
    }
    return sum;
}

这个就是求和，如果要计算数组 a 区间 [left,right] 的和，可以像下面这样调用。

public int sumRange(int left, int right) {
    return prefixSum(right + 1) - prefixSum(left);
}

树状数组的求和我们知道了，那么修改呢（这里先讨论单点修改）？如果树状数组的一个节点值被修改了，那么他的父节点值都要改变，如下图所示，当 a[5] 的值修改后， c5 ，c6 以及 c8 都要修改。

alt

如果要更改 c[i] 的值，只需要找到 c[i] 的父节点以及他父节点的父节点，……，一直往上走，直到根节点，全部修改即可。通过图 1-8 我们可以发现， c[i] 的父节点就是 c[i+lowBit(i)] ，所以我们需要以 c[i] 为起点，通过循环不断的往上找父节点然后修改，来看下树状数组的完整代码：

public class TreeArray {

    private int[] a;// 原始数组，有效下标从 0 开始。
    private int[] c;// 树状数组，有效下标从 1 开始。

    public TreeArray(int[] a) {
        this.a = a;
        this.c = new int[a.length + 1];// 创建树状数组。
        for (int i = 0; i < a.length; i++)
            add(i + 1, a[i]);
    }

    // 把数组 c 中 i 位置的元素加上 val。
    private void add(int i, int val) {
        while (i < c.length) {
            c[i] += val;
            i += lowBit(i);// 继续找父节点。
        }
    }

    // 修改树状数组的值。
    public void update(int i, int val) {
        add(i + 1, val - a[i]);
        a[i] = val;
    }

    // 求数组区间[left,right]的值。
    public int sumRange(int left, int right) {
        return prefixSum(right + 1) - prefixSum(left);
    }

    private int lowBit(int x) {
        return x & -x;
    }

    // 求数组区间[0,i-1]的值。
    private int prefixSum(int i) {
        int sum = 0;
        while (i > 0) {
            sum += c[i];
            i -= lowBit(i);
        }
        return sum;
    }
}

区间更新，单点查询

对于数组更新和查找可以分为下面几类：

单点更新，单点查询
单点更新，区间查询
区间更新，单点查询
区间更新，区间查询

对于“单点更新，单点查询”，原始数组就可以做，不需要构建树状数组。而“单点更新，区间查询”我们上面刚介绍过，这里来看下“区间更新，单点查询”。如果想要区间更新需要构建差分数组，在前面差分数组中我们讲过，如果要更新原始数组区间的值，只需要更新差分数组两边的值即可。其中差分数组的前缀和就是数组 a 中某个元素的值，来看下代码：

public class TreeArray2 {

    private int[] a;// 原始数组，有效下标从 0 开始。
    private int[] d;// 差分数组，有效下标从 1 开始。

    public TreeArray2(int[] a) {
        this.a = a;
        this.d = new int[a.length + 1];// 构建差分数组。
        for (int i = 0; i < a.length; i++)
            add(i + 1, i == 0 ? a[0] : a[i] - a[i - 1]);
    }

    // 把数组 d 中 i 位置的元素加上 val。
    private void add(int i, int val) {
        while (i < d.length) {
            d[i] += val;
            i += lowBit(i);// 继续找父节点。
        }
    }

    // 在数组 a 的区间 [i,j] 内加上 val，注意原始数组的
    // 下标是从 0 开始的，差分数组的下标从 1 开始。
    public void update(int i, int j, int val) {
        add(i + 1, val);   // d[i+1]增加 val。
        add(j + 2, -val); // d[j+2]减少 val。
    }

    private int lowBit(int x) {
        return x & -x;
    }

    // 单点查询，求 a[i-1] 的值。
    public int prefixSum(int i) {
        int sum = 0;
        while (i > 0) {
            sum += d[i];
            i -= lowBit(i);
        }
        return sum;
    }
}

区间更新，区间查询

区间更新可以使用差分数组，那么区间查询呢？假设求数组 a 的前 n 项和，我们来看下公式推导，如下图所示。

alt

a 是原数组， d 是差分数组，这里我们需要构建两个树状数组，其中 d1[i] = d[i]，d2[i] = d[i]*(i-1); 。

private int[] a;// 原始数组，有效下标从 0 开始。
private int[] d1;// 有效下标从 1 开始，(d[1] , d[2] , ... , d[n])。
private int[] d2;// 有效下标从 1 开始，(1*d[2] , 2*d[3] , ... , (n-1)*d[n])。

public TreeArray3(int[] a) {
    this.a = a;
    this.d1 = new int[a.length + 1];
    this.d2 = new int[a.length + 1];
    for (int i = 0; i < a.length; i++)
        add(i + 1, i == 0 ? a[0] : a[i] - a[i - 1]);
}

// 把数组 d 中 i 位置的元素加上 val。
private void add(int i, int val) {
    int x = i - 1;
    while (i < d1.length) {
        d1[i] += val;
        d2[i] += val * x;
        i += lowBit(i);// 继续找父节点。
    }
}

// 在数组 a 的区间 [i,j] 内每个元素加上 val，注意原始数组的
// 下标是从 0 开始的，差分数组的下标从 1 开始。
public void update(int i, int j, int val) {
    add(i + 1, val);   // d[i+1]增加 val。
    add(j + 2, -val); // d[j+2]减少 val。
}

private int lowBit(int x) {
    return x & -x;
}

// 求数组 a 区间[left,right]的值。
public int sumRange(int left, int right) {
    return prefixSum(right + 1) - prefixSum(left);
}

private int prefixSum(int i) {
    int sum = 0;
    int x = i;
    while (i > 0) {
        sum += d1[i] * x - d2[i];
        i -= lowBit(i);
    }
    return sum;
}

区间更新和区间查询除了使用树状数组还可以使用线段树，线段树不光能区间求和，还可以求区间最大值，区间最小值，功能要比树状数组更强大。

#数据结构##数据结构和算法##数据结构与算法面试常考题##面试题##树状数组#

常见数据结构介绍文章被收录于专栏

主要介绍一些常见的数据结构，包括数组，滚动数组，差分数组，树状数组，单向链表，双向链表，循环链表，跳表，异或链表，队列，循环队列，双端队列，栈，散列表，堆，二叉树，二叉搜索树，AVL树，红黑树，字典树，哈夫曼树，线段树，笛卡尔树，图的介绍，图的遍历，Dijkstra算法，Bellman-Ford算法，SPFA算法，Floyd算法，Prim算法，Kruskal算法，Boruvka算法，拓扑排序等。