树状数组

数组
滚动数组
差分数组
树状数组

示例练习>>>

单点更新,区间查询

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


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

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 都要修改。


如果要更改 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 项和,我们来看下公式推导,如下图所示。


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算法拓扑排序

经过两年的打磨,我的新作《算法秘籍》已经完成,后面我会通过视频号直播的方式把书上的知识点都讲一遍,有需要的可以点击购买。也可以点击 内容介绍 查看详情。