树状数组

1.树状数组是什么？

一个数据结构，完成以下两个操作：

• 将第i个数加上c
• 快速求前缀和，即任意区间[i,j]的和

使用数组，第一个操作需要$O(1)$，第二个操作需要$O(n)$，两级分化明显

所以有了树状数组这个结构，使得这两个操作既不会太快，也不会太慢，都维持在$O(logn)$的时间

总的来说：树状数组是一个可以支持快速进行单点修改和区间查询的数据结构

一般树状数组的题目，不会难在代码难写，而是难在你不知道要用树状数组来做

2.代码模板

//树状数组长度是固定的，为n+1
//树状数组的下标必须从1开始
int[] tr = new int[n+1];

// 求最低的一位1
int lowbit(int x){
    return x & -x;
}

// 在tr[x]的位置加上c
void add(int x, int c){
    for(int i = x; i <= n; i += lowbit(i)) tr[i] += c;
}

// 查询前缀和
int query(int x){
    int res = 0;
    for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
    return res;
}

• tr[x]：保存以x为根的子树中叶节点的值的和
• tr[x]节点的长度等于lowbit(x)
• tr[x]节点的父节点为tr[x + lowbit(x)]

3.参考资料(内附有图解)

完全理解并深入应用树状数组

例题1.计算右侧小于当前元素的个数

• 链接：LeetCode315.计算右侧小于当前元素的个数

思路

这是一道非常经典的可以用树状数组处理的问题：逆序对问题

我们可以抽象出一个数组a，a[i]表示数字i的出现次数
我们维护一个树状数组tr，对于a数组中的数据，动态维护a的前缀和，并可以进行单点修改
nums[i]的数据范围在[-10000,10000]
由于tr数组下标从1开始，所以我们需要把这个区间映射到[1,20001]，每个数加10001即可

nums = [5,2,6,1]，映射为[10006,10003,10007,10002]
我们从后往前遍历数组：
第1个数为10002，进行query(10001)，查询1到10001之间所有数的出现次数和，为0，记录答案，更新tr数组
第2个数为10007，进行query(10006)，查询1到10006之间所有数出现的次数和，为1，记录答案，更新tr数组
第3个数为10003，进行query(10002)，查询1到10002之间所有数出现的次数和，为1，记录答案，更新tr数组
第4个数为10006，进行query(10005)，查询1到10005之间所有数出现的次数和，为2，记录答案，更新tr数组

Java代码

class Solution {
    int n = 20001;
    int[] tr = new int[n+1];

    int lowbit(int x){
        return x & -x;
    }

    void add(int x, int c){
        for(int i = x; i <= n; i += lowbit(i)) tr[i] += c;
    }

    int query(int x){
        int res = 0;
        for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }

    public List<Integer> countSmaller(int[] nums) {
        LinkedList<Integer> res = new LinkedList<>();
        for(int i = nums.length-1; i >= 0; i--){
            nums[i] += 10001;
            res.addFirst(query(nums[i]-1));
            add(nums[i],1);
        }
        return res;
    }
}

例题2.数组中的逆序对

• 链接：剑指offer51.数组中的逆序对

思路

这题其实就是上题换了一个形式问，每次把答案累加起来即可

但是有一点不同，这题没有给我们数组元素的范围！
这时候就需要将数组中的元素全部离散化，离散化：排序+去重
我们可以使用TreeSet和HashMap来对元素进行离散化

例如nums=[7,5,6,4]，先用TreeSet进行去重和排序[4,5,6,7]
然后用哈希表进行下标的映射：[4->1, 5->2, 6->3, 7->4]

抽象一个数组a，a[i]表示数字i在哈希表中对应的值所出现的次数
然后就可以用树状数组来模拟了

我们从后往前遍历数组：
第1个数为4，对应下标为1，query(1-1)表示查询下标0(无效下标)对应的值的出现次数，为0，add(1,1)
第2个数为6，对应下标为3，query(3-1)表示查询下标1到下标2对应的值的出现次数，为1，add(3,1)
第3个数为5，对应下标为2，query(2-1)表示查询下标1到下标1对应的值的出现次数，为1，add(2,1)
第4个数为7，对应下标为4，query(4-1)表示查询下标1到下标4对应的值的出现次数，为3，add(4,1)
把四次查询的结果累加：就是逆序对的数量5

Java代码

class Solution {
    int[] tr;
    int n;
    int lowbit(int x){
        return x & -x;
    }

    void add(int x, int c){
        for(int i = x; i <= n; i += lowbit(i)) tr[i] += c;
    }

    int query(int x){
        int res = 0;
        for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }

    public int reversePairs(int[] nums) {
        //离散化
        TreeSet<Integer> set = new TreeSet<>();
        for(int num: nums) set.add(num);

        Map<Integer,Integer> map = new HashMap<>();
        int idx = 1;
        for(int x: set){
            map.put(x, idx++);
        }

        n = set.size();
        tr = new int[n+1];

        int res = 0;
        for(int i = nums.length-1; i >= 0; i--){
            res += query(map.get(nums[i])-1);
            add(map.get(nums[i]),1);
        }

        return res;
    }
}

例题3.翻转对

• 链接：LeetCode493.翻转对

思路

这题又是上一个题目的扩展，先离散化，但是需要将所有nums[i],nums[i]*2的值进行离散化
我们查询的时候，是对nums[i]离散化后的下标进行查询
我们存储的时候，是对nums[i]*2离散化后的下标进行存储，方便前面的元素进行查询

注意：int型的数乘2后可能会爆int，哈希表需要用long类型来存储待离散化的数

具体样例的分析，跟上一题一样，不再赘述

Java代码

class Solution {
    int n;
    int[] tr;
    int lowbit(int x){
        return x & -x;
    }

    void add(int x, int c){
        for(int i = x; i <= n; i += lowbit(i)) tr[i] += c;
    }

    int query(int x){
        int res = 0;
        for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }


    public int reversePairs(int[] nums) {
        TreeSet<Long> set = new TreeSet<>();
        for(int x: nums) {
            set.add((long)x);
            set.add((long)x*2);
        }

        Map<Long,Integer> map = new HashMap<>();
        int idx = 1;
        for(long x: set) map.put(x, idx++);

        n = set.size();
        tr = new int[n+1];

        int res = 0;
        for(int i = nums.length-1; i >= 0; i--){
            res += query(map.get((long)nums[i])-1);
            add(map.get((long)nums[i]*2),1);
        }
        return res;
    }
}