高级数据结构1—初识树状数组—快速求得前缀和和修改某一元素值

- 本人的LeetCode账号：魔术师的徒弟，欢迎关注获取每日一题题解，快来一起刷题呀~

• 本人Gitee账号：路由器，欢迎关注获取博客内容源码。

  树状数组和其他的高级数据结构不同，它非常的好写，同时解决问题也比较局限，所以树状数组的题目的难度主要集中在思考而非代码。

一、基本原理

  树状数组可以解决两个操作：快速的求前缀和、修改某一个数，这两个操作都是O(logn)的。

  这两个操作如果我们直接来操作：

• 存原数组，前缀和O(N)，修改一个数O(1)
• 维护前缀和，前缀和O(1)，修改一个数O(N)

  有一种鱼和熊掌不可兼得的感觉，但是我们的题目中时间复杂度一般取决于最糟糕的时间复杂度，所以如果有n次查询，那么复杂度会达到O(n^2)。树状数组有一个折中的思想，它让这两个操作的时间复杂度都变成了O(logn)，这样总时间复杂度就是O(nlogn)，就会快很多了。

  它是一种基于二进制的方法来解决这个问题的。

  假设我们有一个数x，其二进制表示为：
$$\begin{gathered} x=2^{i_k}+2^{i_{k-1}}+...+2^{i_1} \\ {i}_k>=i_{k-1}>=...>=i_1 \end{gathered}$$
  假设我们想求的是下标为1~x的总和，那么我们可以把1~x这个区间划分成k部分：
$$\begin{gathered} (x-2^{i_1},x] \\ (x-2^{i_1}-2^{i_2},x-2^{i_1}] \\ ... \\ (0,x-2^{i_1}-2^{i_2}-...-2^{i_{k-1}}] \end{gathered}$$
  这样就把下标为1~x这个区间划分成了logx份，这样如果算1~x的总和，只需要求logx个区间的和就能算出来了。

  这个思想就是让我们在logn的时间复杂度中使用前缀和的思想。

  下面来看看区间中元素的个数和区间右端点有什么关系：
$$\begin{gathered} (x-2^{i_1},x],元素个数2_1^i个 \\ (x-2^{i_1}-2^{i_2},x-2^{i_1}]，元素个数2_2^i个 \\ ... \\ (0,x-2^{i_1}-2^{i_2}-...-2^{i_{k-1}}],元素个数2_k^i个 \end{gathered}$$
  区间中元素的个数就是右端点二进制的最低位1的所对应的2的幂。

  所以每个区间可以这样表示：[R - lowbit(R) + 1, R]，其中lowbit(R) = R & (-R)，因此我们可以用一个参数的函数来表示这个区间的总和：C[R]。

  C[x]表示原数组的a[x - lowbit(x) + 1, x]的区间和，所以对原数组下标1~x的和，我们可以用logx个C[x]的和就能求出来了。

  再考虑一下C[x]的关系。

  我们发现，父节点和子结点的关系：
$$\begin{gathered} c[x]=a[x]+c[x-1]+c[x-1-lowbit]+...+c[0](每次都从x-1去掉最后一个1) \\ 14:a[14]+c[13=(01101{)}_2]+c[(01100{)}_2](到0了不算了) \\ 13:a[13]+c[01100](到0了不算了) \\ 16:a[16]+c[15]+c[14]+c[12]+c[8]+c[0] \end{gathered}$$
  如何通过子结点找父节点呢？（对应修改操作）

  修改一个值后，如何确定它会影响哪些父节点呢？

  其实就是上面操作的逆操作，对于一个x，找到其最后一段这样的东西01...10...0,让它进位成10...0即可，发现只要加上10...0（最低位的1的2的幂）即可，这就找到了它的直接父节点，即操作就是x + lowbit(x)。

  修改操作add(int x, int c)：对原数组a，a[x4] += c，就是我们刚刚得到的它会影响哪些父节点，对每个数从0开始的初始修改其实就是创建树状数组。

for (int i = x; i <= n; i += lowbit(i)) tr[i] += c;

  查询操作sum(int x)：求1~x的和，就是我们最早的C[i]的区间含义，tr[i]表示以下标i结尾的，长度为lowbit(i)的区间。

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

  树状数组常用于快速求得前缀和和修改元素组的值、统计一个数组左边或者右边比当前值大或小的元素个数。

  以Less[i] = [0,i)比nums[i]小的元素的个数的求解为例（假设nums元素全部为正），首先我们建立一个树状数组B，大小是nums的最大元素 + 1，然后从左往右遍历，每次先求一下B.sum(nums[i] - 1)，它就是下标1~nums[i] - 1的和，而我们tr[i] = c表示i的元素出现了c次，所以其前缀和sum(x)就是小于等于x的元素出现的次数。

二、例题

1 楼兰图腾

  题意，给了平面上的n个点，它们的横坐标是1~n的一个有序排列，纵坐标是1~n的一个任意排列，统计一下有多少个三元组(i, j ,k),满足两边都比中间高i < j < k && yi > yj && yk > yj，这是第一问，第二问就是中间比两边高。

  数据范围是20w，意味着我们要用一个nlogn的算法。

  首先我们从一个集合考虑，假设一个集合中是所有的点，我们把它以横坐标为1 2 …n分为n个集合，这个划分是不重不漏的。

  那么我们就看看第k部分的满足条件的子集有多少个即可，我们只要统计出yk左边有多少点大于yk，yk右边有多少点大于yk，利用乘法原理两个乘起来就行。

  我们可以从左到右扫描一遍，得到greater[k]表示1~k - 1有多少点的纵坐标大于yk，再从右往左遍历一遍，得到k + 1~n中有多少数大于yk。

  统计一个区间的和，这是树状数组可以解决的问题。

  然后k统计完后，我们就给yk加1即可，这对应修改操作。

  上面写的很乱，建议思路看代码。

#include <iostream>
#include <cstring>
using namespace std;
typedef long long LL;
const int N = 200010;

int n;
int a[N];// 表示点(i, a[i])
int Greater[N];// Greater[x]存当前比x大的数的个数
int Lower[N]; // Lower[x]存当前比x小的数的个数

int tr[N]; // 树状数组对应的原数组含义就是数字i出现了a[i]次
// 因此其前缀和sum(i - 1)就表示小于i的数字出现了多少次
// 从左往右遍历时 它的sum表示当前点左边 1~x中的数字和
// 因为我们每次遍历完一个点就会add(y, 1)
// 所以其实sum(y - 1)表示当前点左边比y小的数字出现的次数
// 同理 从右往左遍历是 sum(y - 1)表示当前点右边出现过的比y小的数字

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

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

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

int main()
{
    scanf("%d", &n);
    for (int i = 1; i <= n; ++i) scanf("%d", &a[i]);
    // 从左往右遍历 确定Greater和Lower
    for (int i = 1; i <= n; ++i)
    {
        int y = a[i];
        // 当前区间左边比y小的数的出现次数
        Lower[i] = sum(y - 1);
        // 当前区间左边比y大的数的出现次数 最大的数是n
        Greater[i] = sum(n) - sum(y);
        // 这个点出现过了 给树状数组加上
        add(y, 1);
    }
    memset(tr, 0, sizeof(tr));// 清空树状数组
    LL res1 = 0;// V的个数
    LL res2 = 0;// ^的个数
    for (int i = n; i >= 1; i--)
    {
        int y = a[i];
        // 统计当前点右边比y大的数的个数
        res1 += (LL)Greater[i] * (sum(n) - sum(y));
        // 统计当前点右边比y小的数字的个数
        res2 += (LL)Lower[i] * sum(y - 1);
        add(y, 1);
    }
    cout << res1 << ' ' << res2 << endl;
    return 0;
}

2 树状数组板子

class BIdxT
{
public:
    BIdxT(int sz)
        : tr(sz + 1)
    {}
    int lowbit(int x) { return x & (-x); }
    void add(int k, int val)
    {
        for (int i = k; i < tr.size(); i += lowbit(i)) tr[i] += val;
    }
    void reinit()
    {
        fill(tr.begin(), tr.end(), 0);
    }
    int sum(int x)
    {
        int res = 0;
        for (int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }
private:
    vector<int> tr;
};

3 LeetCode307. 区域和检索 - 数组可修改

307. 区域和检索 - 数组可修改

  本题需要的两个操作是快速求得前缀和，快速修改一个数组中的值，显然可以用树状数组来处理。

class tr
{
public:
    tr(const vector<int>& nums)
        :t(nums.size() + 1)
    {
        for (int i = 0; i < nums.size(); ++i) add(i + 1, nums[i]);
    }
    int lowbit(int x)
    {
        return x & (-x);
    }
    void add(int x, int val)
    {
        for (int i = x; i < t.size(); i += lowbit(i)) t[i] += val;
    }
    int sum(int x)
    {
        int res = 0;
        for (int i = x; i > 0; i -= lowbit(i)) res += t[i];
        return res;
    }
    int query(int left, int right)
    {
        return sum(right) - sum(left - 1);
    }
private:
    vector<int> t;
};

class NumArray {
public:
    NumArray(vector<int>& nums) 
        : _tr(nums), num(nums)
    {}
    
    void update(int index, int val) 
    {
        _tr.add(index + 1, val - num[index]);
        num[index] = val;
    }
    
    int sumRange(int left, int right) 
    {
        return _tr.query(left + 1, right + 1);
    }
private:
    tr _tr;
    vector<int>& num;
};

4 LeetCode327. 区间和的个数

327. 区间和的个数

  本题的关键在于意识到要找的满足数量的s(i, j)等价于
$$\begin{gathered} lower<=preSum[j]-presum[i]<=upper \\ preSum[j]-upper<=preSum[i]<=preSum[j]-lower,0<=i<j \\ 就是找[0,j)区间内满足preSum[i]属于上面那个范围的i的数量 \\ 可以用树状数组维护，从左向右遍历时，得到小于preSum[j]-upper和小于preSum[j]-lower的数量 \\ 作差即得到当前满足条件的数量，累计求和即可 \end{gathered}$$

  本题的前缀和得到的数会比较大，用LL存一下，由于数据个数其实不多而且有负数，所以可以做一个离散化处理。

typedef long long LL;
// 树状数组
class BIndexT
{
public:
    BIndexT(int sz)
        : tr(sz + 1)
    {}
    int lowbit(int x)
    {
        return x & (-x);
    }
    void add(int k, int val)
    {
        for (int i = k; i < tr.size(); i += lowbit(i)) tr[i] += val;
    }
    int sum(int x)
    {
        int res = 0;
        for (int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }
private:
    vector<int> tr;
};
class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) 
    {
        vector<LL> preSum(nums.size() + 1);
        // 求解前缀和
        for (int i = 0; i < nums.size(); ++i)
        {
            preSum[i + 1] = preSum[i] + nums[i];
        }
        vector<LL> alls;
        for (LL p : preSum)
        {
            alls.push_back(p);
            alls.push_back(p - upper);
            alls.push_back(p - lower);
        }
        // 离散化
        sort(alls.begin(), alls.end());
        alls.erase(unique(alls.begin(), alls.end()), alls.end());
        unordered_map<LL, int> fix;
        int idx = 0;
        for (LL num : alls)
        {
            fix[num] = idx++;
        }
        // 创建树状数组 其sum表示小于x的元素出现的次数
        BIndexT B(alls.size());
        int res = 0;
        // 遍历前缀和数组
        // 查询前缀和下标(0, i)中值处于区间[preSum[i] - upper, preSum[i] - lower]的个数
        for (int i = 0; i < preSum.size(); ++i)
        {
            LL p = preSum[i];
            int l = fix[p - upper];
            int r = fix[p - lower];
            int L = B.sum(l);
            int R = B.sum(r + 1);
            res += R - L;
            B.add(fix[p] + 1, 1);
        }
        return res;
    }
};

5 LeetCode.1395统计作战单位数

1395. 统计作战单位数

  本题也是经典的树状数组应用题：统计左边有多少元素比自己小Less[i]，统计左边有多少元素比自己大Greater[i]，统计右边有多少元素比自己小，统计右边有多少元素比自己大。

  因为士兵得分都是正值，所以不必离散化直接干就完了。

class BIdxT
{
public:
    BIdxT(int sz)
        : tr(sz + 1)
    {}
    int lowbit(int x) { return x & (-x); }
    void add(int k, int val)
    {
        for (int i = k; i < tr.size(); i += lowbit(i)) tr[i] += val;
    }
    void reinit()
    {
        fill(tr.begin(), tr.end(), 0);
    }
    int sum(int x)
    {
        int res = 0;
        for (int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }
private:
    vector<int> tr;
};
class Solution {
public:
    int numTeams(vector<int>& rating) 
    {
        vector<int> Greater(rating.size());
        vector<int> Lower(rating.size());
        int n = *max_element(rating.begin(), rating.end());
        BIdxT B(n);
        for (int i = 0; i < rating.size(); ++i)
        {
            int score = rating[i];
            Greater[i] = B.sum(n) - B.sum(score);
            Lower[i] = B.sum(score - 1);
            B.add(score, 1);
        }
        B.reinit();
        int res1 = 0;// > >
        int res2 = 0;// < <
        for (int i = rating.size() - 1; i >= 0; --i)
        {
            int score = rating[i];
            res1 += Greater[i] * B.sum(score - 1);
            res2 += Lower[i] * (B.sum(n) - B.sum(score));
            B.add(score, 1);
        }
        return res1 + res2;
    }
};

6 LeetCode 315. 计算右侧小于当前元素的个数

315. 计算右侧小于当前元素的个数

  本题也是比较典型的寻找该位置右边比自己小的元素的个数，由于有负数，所以要离散化处理一下。

class BIdxT
{
public:
    BIdxT(int sz)
        : tr(sz + 1)
    {}
    int lowbit(int x) { return x & (-x); }
    void add(int k, int val)
    {
        for (int i = k; i < tr.size(); i += lowbit(i)) tr[i] += val;
    }
    void reinit()
    {
        fill(tr.begin(), tr.end(), 0);
    }
    int sum(int x)
    {
        int res = 0;
        for (int i = x; i > 0; i -= lowbit(i)) res += tr[i];
        return res;
    }
private:
    vector<int> tr;
};
class Solution {
public:
    vector<int> countSmaller(vector<int>& nums) 
    {
        // 离散化
        vector<int> alls(nums);
        sort(alls.begin(), alls.end());
        alls.erase(unique(alls.begin(), alls.end()), alls.end());
        unordered_map<int, int> myhash;
        int idx = 1;
        for (int p : alls)
        {
            myhash[p] = idx++;
        }
        BIdxT B(alls.size());
        vector<int> Lower(nums.size());
        for (int i = nums.size() - 1; i >= 0; --i)
        {
            int num = nums[i];
            Lower[i] = B.sum(myhash[num] - 1);
            B.add(myhash[num], 1);
        }
        return Lower;
    }
};