[树状数组][前缀和] Jzoj P4738 神在夏至祭降下了神谕

Description

 

Input

Output

 

Sample Input

4 1
0 0 1 1

Sample Output

5

 

Data Constraint

 

Hint

 

 

题解

• 设f[i]为前i个人的划分方案数
• f[i]=∑f[j](|sum[i]-sum[j]|<=k)
• sum可以用前缀和预处理
• 发现，如果sum[j]在[sum[i]-k,sum[i]+k]中就可以转移
• 那么就可以用树状数组来做，然后sum[i]可能为负数，c++又没有负数数组，所以要开两倍长度

代码

#include <cstdio>
#include <iostream>
using namespace std;
int n,k,sum[100010],x;
long long sz[400010],f[100010],mo=1e9+7;
void insert(int x,int y) { for (x+=200000;x<=400010;x+=x&-x) sz[x]+=y; }
long long getsum(int x)
{
    long long y=0;
    for (x+=200000;x;x-=x&-x) y=(y+sz[x])%mo;
    return y;
}
int main()
{
    scanf("%d%d",&n,&k);
    for (int i=1;i<=n;i++)
    {
        scanf("%d",&x);
        sum[i]=x?sum[i-1]+1:sum[i-1]-1;
    }
    insert(0,1);
    for (int i=1;i<=n;i++) f[i]=(getsum(sum[i]+k)-getsum(sum[i]-k-1)+mo)%mo,insert(sum[i],f[i]);
    printf("%lld",f[n]);
    return 0;
}