火速上手树状数组

2022-03-08 算法

lowbit
树状数组区间查询
树状数组单点修改
树状数组区间修改

树状数组的代码非常简单，但原理十分复杂，还涉及到一些不算容易的前置知识。这里总结出树状数组的用法与代码，旨在帮助初学者快速上手树状数组这一强大的数据结构。

lowbit

lowbit(x)指一个整数x去掉二进制表示下所有高位的1之后剩下的数。

比如，10（十）的二进制表示是1010，lowbit(10)就是0010，即lowbit(10) = 2 。

lowbit(x)的代码如下。

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

后续的代码都需要用到这里的lowbit。

树状数组区间查询

树状数组能快速回答任意区间的和的询问，时间复杂度 $O(logn)$ 。用树状数组本身能得到的的只是前缀和，求区间的的值需要用两个前缀和做差（当然，如果你不理解什么是前缀和也没有关系，会写代码就好了）。

int sum(int x, int *bit) //求x处的前缀和
{
    int ret = 0;
    while (x > 0)
    {
        ret += bit[x];
        x -= lowbit(x);
    }
    return ret;
}
int query(int l, int r) //求闭区间[l,r]所有元素的和
{
    return sum(r) - sum(l - 1);
}

树状数组单点修改

在树状数组上修改一个点的值的时间复杂度是 $O(logn)$ ，代码如下

void point_add(int num, int x, int *bit) //将序列的第num个元素加上x
{
    while (num <= n) //n:序列的总长度
    {
        bit[num] += x;
        num += lowbit(num);
    }
}

上面的这个函数实现的功能是将序列的第num个元素加上x。如果是想“把序列的第num个元素变成x”，就相当于

“将序列的第num个元素加上x-a”，其中a为该元素本来的值。

树状数组不能直接访问到一个元素的值，所以不能直接得到上面的所说的a。因此如果需要“把序列的第num个元素变成x”需要先调用之前的区间查询函数得到a，再“将序列的第num个元素加上x-a”

void turn(int num, int x, int *bit) //将第i个元素的值变成x
{
    int a = query(num, num);
    point_add(num, x - a, bit);
}

树状数组区间修改

两个树状数组联动可以通过4次单点修改实现区间修改。这其中用到了差分的思想。

//将闭区间[l,r]上所有元素加上x
void internal_add(int l, int r, int x, int *bit1, int *bit2)
{
    point_add(l, -x * (l - 1), bit1);
    point_add(l, x, bit2);
    point_add(r + 1, x * r, bit1);
    point_add(r + 1, -x, bit2);
}

由于变成了两个树状数组联动，如果需要使用区间修改功能的话，之前的单点修改就不能直接用了，也需要用这个函数来进行单点修改（使区间左端点等于右端点）。同时，之前的查询也不能用了，需要用到下面这个函数

//求闭区间[l,r]所有元素的和
int internal_query(long long l, long long r, long long *bit1, long long *bit2)
{
    return sum(r, bit1) + sum(r, bit2) * r - sum(l - 1, bit1) - sum(l - 1, bit2) * (l - 1);
}

使用例：POJ3468，区间修改模板题

#include <bits/stdc++.h>
using namespace std;
long long b1[500005], b2[500005], n;
long long lowbit(long long x)
{
    return x & -x;
}
long long sum(long long x, long long *bit) //求x处的前缀和
{
    long long ret = 0;
    while (x > 0)
    {
        ret += bit[x];
        x -= lowbit(x);
    }
    return ret;
}
void point_add(long long num, long long x, long long *bit)
{
    while (num <= n)
    {
        bit[num] += x;
        num += lowbit(num);
    }
}
void internal_add(long long l, long long r, long long x, long long *bit1, long long *bit2)
{
    point_add(l, -x * (l - 1), bit1);
    point_add(l, x, bit2);
    point_add(r + 1, x * r, bit1);
    point_add(r + 1, -x, bit2);
}
long long internal_query(long long l, long long r, long long *bit1, long long *bit2)
{
    return sum(r, bit1) + sum(r, bit2) * r - sum(l - 1, bit1) - sum(l - 1, bit2) * (l - 1);
}
int main()
{
    long long q, l, r, a, i;
    char o[5];
    scanf("%lld%lld", &n, &q);
    for (i = 1; i <= n; i++)
    {
        scanf("%lld", &a);
        internal_add(i, i, a, b1, b2);
    }
    while (q--)
    {
        scanf("%s", o);
        if (o[0] == 'Q')
        {
            scanf("%lld%lld", &l, &r);
            printf("%lld\n", internal_query(l, r, b1, b2));
        }
        else
        {
            scanf("%lld%lld%lld", &l, &r, &a);
            internal_add(l, r, a, b1, b2);
        }
    }
    return 0;
}