数据结构学习NOTE - 树状数组

树状数组

高效维护前缀和，加上差分可解决许多问题。

P3374 【模板】树状数组 1

先看如何求 $a[\mathrm{1...7}]$ 的前缀和。

一种做法是直接求 $a_1+a_2+a_3+a_4+a_5+a_6+a_7$，但如果已知 $A=a[\mathrm{1...4}],B=a[\mathrm{5...6}],C=a[\mathrm{7...7}]$ ，当然只用将三个数相加就行。

树状数组就凭借将一段前缀 $[\mathrm{1...}n]$ 拆成长度不超过 $\log (n)$ 个区间，使这些区间信息已知。

详细学习参照 树状数组。

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>

using namespace std;

const int maxn=5e5+5;
int a[maxn],tr[maxn];
int n,m,x,y,opt;

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

void add(int x,int k)
{
	while(x<=n) tr[x]+=k,x+=lowbit(x);
}

int query(int x)
{
	int sum=0;
	while(x>0) sum+=tr[x],x-=lowbit(x);
	return sum;
}

int main()
{
	#ifndef ONLINE_JUDGE
	//freopen("in.txt","r",stdin);
	#endif
	cin>>n>>m;
	for(int i=1;i<=n;i++) scanf("%d",&a[i]),add(i,a[i]);
	for(int i=1;i<=m;i++)
	{
		scanf("%d%d%d",&opt,&x,&y);
		if(opt==1) add(x,y);
		if(opt==2)printf("%d\n",query(y)-query(x-1));
	}
	return 0;
}

P3368 【模板】树状数组 2

上面的问题是单点修改+区间查询，这里是区间修改+单点查询。由查分数组定义：${\sum }_{x=1}^{k}d[i]=a[k]$ ，这里考虑使用树状数组维护差分数组，即可在 $O(\log n)$ 的时间做区间修改。

而对于单点查询，因为差分是前缀和的逆运算，即查分数组的前缀和数组是原数组，即求 $1$ 至 $k$ 的前缀和就是答案。

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>

using namespace std;

const int maxn=500005;
int tr[maxn];
int n,m,x,y,k,lst=0;
int lowbit(int x)
{
	return x&-x;
}

void add(int x,int k)
{
	while(x<=n) tr[x]+=k,x+=lowbit(x);
}

int query(int x)
{
	int sum=0;
	while(x>0) sum+=tr[x],x-=lowbit(x);
	return sum;
}

int main()
{
	#ifndef ONLINE_JUDGE
	//freopen("in.txt","r",stdin);
	#endif
	cin>>n>>m;
	
	for(int i=1;i<=n;i++)
	{
		int t;
		scanf("%d",&t);
		add(i,t-lst);
		lst=t;
	}
	for(int i=1;i<=m;i++)
	{
		int opt;
		cin>>opt;
		if(opt==1)
		{
			scanf("%d%d%d",&x,&y,&k);
			add(x,k);
			add(y+1,-k);
		}
		else
		{
			scanf("%d",&x);
			printf("%d\n",query(x));
		}
	}

	return 0;
}

P3372 【模板】线段树 1

使用树状数组解决区间修改+区间查询问题，使用到二阶树状数组，保留上一道题的差分思想，我们又有一下式子：

$$\begin{array}{rl}\sum _{i=1}^k{a}_i& =\sum _{i=1}^k\sum _{j=1}^i{D}_j\\ & =\sum _{i=1}^k(k-i+1)D_i\\ & =k\sum _{i=1}^k{D}_i-\sum _{i=2}^{k-1}(i-1)D_i\\ & =k\sum _{i=1}^k{D}_i-\sum _{i=1}^k{D}_i\end{array}$$

即：

$$\begin{array}{rl}{a}_1+a_2+a_3+...+a_k& =D_1+(D_1+D_2)+(D_1+D_2+D_3)+...+(D_1+D_2+D_3+...+D_k)\\ & =k{D}_1+(k-1)D_2+(k-2)D_3+...+(k-(k-1))D_k\\ & =k(D_1+D_2+D_3+...+D_k)-(D_2+2D_3+3D_4+...+(k-1)D_k)\\ & =k\sum _{i=1}^k{D}_i-\sum _{i=1}^k(i-1)D_i\end{array}$$

观察到公式最后一行是求两个前缀和，用两个树状数组维护，一个实现 $D_i$，一个实现 $(i-1)D_i$。

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>
#define int long long

using namespace std;

const int maxn=1e5+5;
int tr[3][maxn],a[maxn];
int n,m,opt,x,y,k;

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

void add(int x,int k,int t)
{
	while(x<=n)
	{
		tr[t][x]+=k;
		x+=lowbit(x);
	}
}

int query(int x,int t)
{
	int sum=0;
	while(x>0)
	{
		sum+=tr[t][x];
		x-=lowbit(x);
	}
	return sum;
}

signed main()
{
	#ifndef ONLINE_JUDGE
	//freopen("in.txt","r",stdin);
	#endif
	cin>>n>>m;
	for(int i=1;i<=n;i++)
	{
		scanf("%lld",&a[i]);
	}
	for(int i=1;i<=n;i++)
	{
		add(i,a[i]-a[i-1],1);
		add(i,(i-1)*(a[i]-a[i-1]),2);
	}
	while(m--)
	{
		scanf("%d",&opt);
		if(opt==1)
		{
			scanf("%lld%lld%lld",&x,&y,&k);
			add(x,k,1),add(y+1,-k,1);
			add(x,k*(x-1),2),add(y+1,-k*y,2);
		}
		else
		{
			scanf("%lld%lld",&x,&y);
			printf("%lld\n",y*query(y,1)-query(y,2)-(x-1)*query(x-1,1)+query(x-1,2));
		}
	}
	return 0;
}

244. 谜一样的牛

非常好的一道题，使用树状数组+二分将复杂度降至 $O(n\log (n{)}^2)$ 。

• 题目大意：

  给定每头牛前面的牛身高低于它的牛有多少，保证每头牛的身高为 $1\sim n$ 的一个排列，求出每头牛的身高。

• 题目分析：

  先看样例是如何操作的：

  1. 对于第5头奶牛，此时可供选择的身高有 1 2 3 4 5，因为前面没有比他矮的奶牛，它又是最后一头，故身高只能是 $1$。

  2. 对于第4头奶牛，此时可供选择的身高有 2 3 4 5，由前面有一头身高比他矮的牛，故他只能取到身高序列中的次小值，为$3$。

  3. 对于第三头奶牛，此时可供选择身高有 2 4 5，前面有两头比他矮的牛，它的身高只能为第三小值 $5$。

  依次类推，于是我们可以发现从后往前枚举，第 $i$ 头奶牛的身高为可选身高序列中第 $a_i+1$ 大值。

  将问题转化为两个操作：

  1. 删除一个身高。
  2. 求出第 $k$ 小身高。

  假定维护一个数组 $a$ ，一开始 $a$ 中所有元素为 $1$ ，表示身高 $i$ 还没有被选；如果被选了，则 $a_i=0$ 。那么删除操作就好做了。

  接下来考虑求第 $k$ 小数，由于 $a$ 中元素只有 $0,1$ 两种可能，那么 $a$ 的前缀和肯定是 不降 的，满足二分条件，直接二分即可。

  求第 $k$ 小数 $\Leftrightarrow $ 找出前缀和（$1\sim n$ ）大于等于 $k$ 的第一个位置。

  此时复杂度 $O(n^2)$ 不足以通过此题。

  又观察到关于数组 $a$ 操作的性质：只用支持求前缀和+单点修改即可。

  考虑使用树状数组。可将复杂度降至 $O(n\log (n{)}^2)$ 。

• AC_code

  代码如下：

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>

using namespace std;

const int maxn=1e5+5;
int n;
int tr[maxn],a[maxn],num[maxn],ans[maxn];
inline int lowbit(int x)
{
	return x&-x;
}
void add(int x,int k)
{
	for(int i=x;i<=n;i+=lowbit(i)) tr[i]+=k; 
}
int query(int x)
{
	int sum=0;
	for(int i=x;i>0;i-=lowbit(i)) sum+=tr[i];
	return sum;
}

int main()
{
	#ifndef ONLINE_JUDGE
	//freopen("in.txt","r",stdin);
	#endif
	cin>>n;
	for(int i=2;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1;i<=n;i++) add(i,1),ans[i]=0;
	for(int i=n;i>=1;i--)
	{
		//找到剩余数中第k大
		int l=1,r=n;
		while(l<=r)
		{
			int mid=(l+r)/2;
			if(query(mid)>=a[i]+1)
			{
				r=mid-1;
			}
			else
			{
				l=mid+1;
			}
//			cout<<l<<" "<<r<<endl;
		}
		ans[i]=l;
		add(l,-1);
	}
	for(int i=1;i<=n;i++) cout<<ans[i]<<endl;
	return 0;
}