稀有猿诉

十年磨一剑,历炼出锋芒,说话千百句,不如码二行。

树状树组简介

树状数组,即Binary Indexed Tree,简单来理解就是用数组来表示一颗树,它的实际存储结构是数组,但元素之间的逻辑关系是树。通常用于解决区间问题和快速计算前缀和的问题。

区间查询和区间修改

在介绍树状数组之前,先来看一个问题,给定数组nums,长度为n,如何快速实现如下四种操作:

操作 参数 说明
单点查询 i 查询索引为i的元素值
区间查询 l, r 查询区间[l, r]内元素之和
单点修改 i 修改索引 为i的元素值
区间修改 l, r 修改区间[l, r]内的元素值,比如都加上一个数x

首先,肯定 可以想到朴素做法,简单的优化就是用前缀和(参考此文),可以做如下对比:

操作 方法 时间复杂度 方法 时间复杂度 方法 时间复杂度
单点查询 朴素做法 O(1) 前缀和 O(1) 差分数组 O(n)
区间查询 朴素做法 O(n) 前缀和 O(1) 差分数组 O(n)
单点修改 朴素做法 O(1) 前缀和 O(n) 差分数组 O(1)
区间修改 朴素做法 O(n) 前缀和 O(n) 差分数组 O(1)

可以看到基本的方法必然会有一种操作会达到O(n)复杂度。而用树状数组就可以区间相关的操作(区间查询和区间修改)都降到O(logn)。

注意:单点操作,可以视为区间操作的特例,比如区间查询i,可以视为区间查询[i, i]。

树状数组的思想

树状数组就是把原数组分成多个区间,用一个辅助数组来存储这些划分出来的区间的区间和。普通前缀和,对于区间查询 和单点查询 来说肯定 可以O(1)来完成,但对于单点修改,比如要修改nums[i]的值,这会影响到包含nums[i]的所有的区间和,因此会是O(n)的复杂度。树状数组因为预先划分了一些区间,能能够保证受nums[i]影响的区间只有O(logn)个,因此单点修改对于树状数组的复杂度只会是O(logn)。

树状数组是一个辅助数组tree[],长度是原数组长度n+1,区间划分的原则是从[1,n]中的每一个下标i,存储的是一个区间和,这个区间是(i-lowbit(i),i],注意区间是左开右闭。这个lowbit是一个函数取的是索引i的最左一个1所代表的数字,比如索引4,其lowbit(4) = (0b100)2=4,所以tree[4]存储的是(0,4]这个区间的区间和;而5,lowbit(5)=(0b101)2=1,所以tree[5]只存它自己。

这样划分区间后,对于单点修改,只需要把包含影响的元素的区间更新一遍就可以了。比如修改元素i,那么只需要把[i, i+lowbit(i)]的元素更新一下就可以了。

对于区间查询,区间查询用前缀和可以做到O(1),因此问题的关键在于利用划分的区间和区间和来找到前缀和。区间划分大小为lowbit(i),tree[i]存储的是(i-lowbit(i), i]的和,那么,如果令j=i-lowbit(i),那么tree[j]就是(j-lowbit(j),j]的和,这样一直到索引0,就找到了到i的前缀和了。由此可以看出用下标索引lowbit来划分区间可以非常巧妙的分治单点修改和区间查询的问题,使得复杂度都降为O(logn)。

lowbit

区间划分最为精妙 的就是这个lowbit,它是能保证tree中所有的查询和修改都能在O(logn)时间内完成。它的作用是取索引二进制形式的最后一个1所代表的数字。它的实现并不复杂,用一点位运算的技巧就可以了,位运算技巧可以参考文章

1
2
3
private static int lowbit(int i) {
    return i & -i;
}

细心的你肯定可以发现,对于奇数索引lowbit肯定是1,所以tree中奇数索引存储的都是原数组的元素,只有偶数索引存储的才是区间之和。

基础树状数组的实现

基于前面的原理,可以实现基础版本的树状数组,基础版本的意思是,对于单点修改和区间查询两个操作可以实现O(logn)复杂度的支持。

实现细节需要注意,树状数组是一个辅助数组tree,它的长度是原数组长度加1,索引为0的是不使用的,仅[1, n]是有效数据。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
/**
 * The Binary Indexed Tree.
 */
public class BIT {
    // The original array, which is not really necessary.
    private int[] nums;

    // the size of the Binary Indexed array, size = n + 1, n is the length of original array.
    private int size;

    // The Binary Indexed Array, valid data are in [1, size), where size = n + 1
    // Original array element [i] is mapped to [i + 1] in the tree.
    private int[] tree;

    public BIT(int[] nums) {
        this.nums = nums;
        size = nums.length + 1;

        tree = new int[size];
        for (int i = 0; i < size - 1; i++) {
            add(i + 1, nums[i]);
        }
    }

    private static int lowbit(int i) {
        return i & -i;
    }

    // Update array element [i] to val
    // Meaning set nums[i] to val, but it is easier to maintain tree with delta
    // nums[i] = val, is the same to nums[i] = nums[i] + (val - nums[i]), where val - nums[i]
    // is the delta.
    // So we add affected items in tree with the delta.
    public void update(int i, int val) {
        add(i + 1, val - nums[i]);
        nums[i] = val;
    }

    private void add(int i, int x) {
        for (int j = i; j < size; j += lowbit(j)) {
            tree[j] += x;
        }
    }

    // Query the regional sum in [l, r]
    public int query(int l, int r) {
        return preSum(r + 1) - preSum(l);
    }

    // Get the preSum up to i inclusive.
    private int preSum(int i) {
        int sum = 0;
        for (int j = i; j > 0; j -= lowbit(j)) {
            sum += tree[j];
        }
        return sum;
    }

    public static void main(String[] args) {
        int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9};
        BIT bit = new BIT(a);

        System.out.println("3 = " + bit.query(0, 1));
        System.out.println("18 = " + bit.query(2, 5));

        bit.update(3, 5);
        System.out.println("3 = " + bit.query(0, 1));
        System.out.println("19 = " + bit.query(2, 5));
    }
}

完整代码看这里

中级树状数组

基础版本只能支持单点修改,区间查询。但如果想要区间修改,单点查询的话使用基础版本就会达到O(nlogn)的复杂度,因为每次单点修改会是O(logn)的,那么区间修改就会达到O(nlogn)。

前缀和与差分数组简介这个文章中介绍过差分数组就是前缀和的逆运算,差分数组支持常数级的区间修改,而对差分数组求前缀和就能得到原数组。基于这个,我们可以把树状数组中的输入数组换成原数组的差分,这样tree就变成了差分的区间和,进尔前缀和就变成了原数组,由此便可以实现对数级别的区间修改和单点查询。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
/**
 * The Binary Indexed Tree
 * Support point query (PQ) and range update (RU) in O(log^n) complexity.
 */
public class PQRUBIT {
    private int[] tree;
    private int size;

    public PQRUBIT(int[] nums) {
        int n = nums.length;
        size = n + 1;
        int[] diff = new int[n];
        tree = new int[size];

        diff[0] = nums[0];
        for (int i = 1; i < n; i++) {
            diff[i] = nums[i] - nums[i - 1];
        }
        for (int i = 0; i < size - 1; i++) {
            add(i + 1, diff[i]);
        }
    }

    private void add(int i, int x) {
        for (int j = i; j < size; j += lowbit(j)) {
            tree[j] += x;
        }
    }

    private int lowbit(int i) {
        return i & -i;
    }

    public int query(int i) {
        int res = 0;
        for (int j = i + 1; j > 0; j -= lowbit(j)) {
            res += tree[j];
        }
        return res;
    }

    public void update(int l, int r, int x) {
        add(l + 1, x);
        add(r + 2, -x);
    }

    public static void main(String[] args) {
        int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9};
        PQRUBIT bit = new PQRUBIT(a);

        System.out.println("1 = " + bit.query(0));
        System.out.println("5 = " + bit.query(4));

        bit.update(2, 5, 3);
        System.out.println("1 = " + bit.query(0));
        System.out.println("6 = " + bit.query(2));
        System.out.println("8 = " + bit.query(4));
        System.out.println("9 = " + bit.query(5));
        System.out.println("7 = " + bit.query(6));
    }
}

完整代码看这里

需要特别注意数据对齐,输入数组和差分数组diff都是从0到n-1的,而区间和数组,也即树状数组是从1到n的。但开放出来的接口肯定 也都是从0到n-1的,需要注意当使用nums和diff时,如果是对tree的操作需要对索引做加1。比如BIT中的query在调用preSum时已把索引加1;和PQRUBIT中的query要从j = i + 1开始。

高级树状数组

最高级的树状数组就是支持O(logn)级别的区改区查,要结合前面两种。这个较为复杂一些,需要做一些推导。

先来看差分数组,详细的可以看文章

1
2
3
4
diff[0] = nums[0]
diff[1] = nums[1] - nums[0]
diff[2] = nums[2] - nums[1]
diff[i] = diff[i] - diff[i - 1]

原数组是差分数组的前缀和:

1
2
3
nums[0] = diff[0]
nums[1] = diff[0] + diff[1]
nums[i] = sum(diff, 0, i)

而原数组的前缀和:

1
preSum[k] = nums[0] + nums[1] + ... nums[k]

把差分数组与原数组的关系代入原数组的前缀和,可得:

1
preSum[k] = diff[0] + (diff[0] + diff[1]) + ... (diff[0] + diff[1] + ... diff[k])

把括号展开,合并同类项,可得:

1
preSum[k] = (k+1) * diff[0] + ((k+1) - 1) * diff[1] + ... ((k + 1) - (k-1)) * diff[k-1] + ((k + 1) - k) * diff[k]

再把,括号展开,以k为中心合并,可得:

1
preSum[k] = (k + 1) * (diff[0] + diff[1] + ... diff[k]) - (diff[1] + 2 * diff[2] + k * diff[k])

从PQRUBIT中可知,对于以差分数组为基础的BIT,其query(k) = nums[k] = sum(diff, 0, k),代入上式,可得:

1
preSum[k] = (k + 1) * query(k) - (0 * diff[0] + 1 * diff[1] + 2 * diff[2] + k * diff[k])

至于,后面这一坨,因为它也直接与差分数组diff产生关系,所以可以引入另外一个辅助数组用以存储,进而就可以求出原数组的区间和,也就解决区间查询的问题。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
/**
 * The Binary Indexed Tree
 * Support range query (RQ) and range update (RU) in O(log^n) complexity.
 */
public class RQRUBIT {
    private int size;
    private int[] tree;
    private int[] conTree;

    public RQRUBIT(int[] nums) {
        final int n = nums.length;
        size = n + 1;

        int[] diff = new int[n];
        diff[0] = nums[0];
        for (int i = 1; i < n; i++) {
            diff[i] = nums[i] - nums[i - 1];
        }

        tree = new int[size];
        for (int i = 0; i < n; i++) {
            add(tree, i + 1, diff[i]);
        }

        conTree = new int[size];
        for (int i = 0; i < n; i++) {
            add(conTree, i + 1, i * diff[i]);
        }
    }

    public int rangeQuery(int l, int r) {
        int preSumL = l * query(tree, l) - query(conTree, l);
        int preSumR = (r + 1) * query(tree, r + 1) - query(conTree, r + 1);
        return preSumR - preSumL;
    }

    public void rangeUpdate(int l, int r, int x) {
        add(tree, l + 1, x);
        add(tree, r + 2, -x);

        add(conTree, l + 1, l * x);
        add(conTree, r + 2, (r + 1) * -x);
    }

    private void add(int[] t, int i, int x) {
        for (int j = i; j < size; j += lowbit(j)) {
            t[j] += x;
        }
    }

    private int query(int[] t, int k) {
        int res = 0;
        for (int j = k; j > 0; j -= lowbit(j)) {
            res += t[j];
        }
        return res;
    }

    private int lowbit(int i) {
        return i & -i;
    }

    public static void main(String[] args) {
        int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9};
        RQRUBIT bit = new RQRUBIT(a);

        System.out.println("3 = " + bit.rangeQuery(0, 1));
        System.out.println("18 = " + bit.rangeQuery(2, 5));

        bit.rangeUpdate(3, 6, 3);
        System.out.println("3 = " + bit.rangeQuery(0, 1));
        System.out.println("27 = " + bit.rangeQuery(2, 5));
        System.out.println("10 = " + bit.rangeQuery(2, 3));
    }
}

完整代码看这里

典型题目

题目 题解 说明
307. 区域和检索 - 数组可修改 题解 单点修改,区间查询,标准的BIT
315. 计算右侧小于当前元素的个数 题解
题解
题解
题解
题解

参考资料

Comments