[树状数组]jzoj4738. 神在夏至祭降下了神谕

Description

Input 

Output 

Sample Input 

4 1
0 0 1 1

Sample Output

5

Data Constraint

Hint 

 Solution

很容易想到用dp，设f[i]为在i处划分的方案数，转移为f[i]+=f[j](j）。

考虑优化，我们用pre[x]表示1~x的夏之军人数减冬之军人数，对于f[i]可以转移的f[j]存在以下关系

pre[i]-pre[j]<=k

pre[i]-pre[j]>=-k

可得

pre[i]-k<=pre[j]<=pre[i]+k

那么就可以用树状数组维护这个区间的f[j]总和，时间O（nlogn）

code

代码敲级丑

#include
#include
#include
using namespace std;
#define N 100000+200
#define lowbit(x) x&(-x)
long long n,m;
long long f[N],tree[4*N],pre[N];
long long ti=200000;
long long mod=1000000007;
void add(long long x,long long ci)
{
	while(x<=300000)
	{
		tree[x]+=ci;
		x+=lowbit(x);
	}
}
long long query(long long x)
{
	long long cnt=0;
	while(x>0)
	{
		cnt+=tree[x];
		x-=lowbit(x);
	}
	return cnt;
}
int main()
{
	scanf("%lld %lld",&n,&m);
	for(long long i=1;i<=n;i++)
	{
		long long x;
		scanf("%lld",&x);
		if(x==1)
			pre[i]=pre[i-1]+1;
		else
			pre[i]=pre[i-1]-1;
	}
	for(long long i=1;i<=n;i++)
	{
		if(abs(pre[i])<=m)
			f[i]++;
		f[i]=(f[i]+query(pre[i]+m+ti)-query(pre[i]-m-1+ti))%mod;
		add(pre[i]+ti,f[i]);
	}
	printf("%lld",f[n]%mod);
}