【挖坑记】JZOJ 4738 神在夏至祭降下了神谕

题目大意

一个长度为n的01序列要分段，每一段的0、1个数不能相差k以上，问有多少种分段方案。
n<=1e5,k<=n
时间限制 1s
空间限制 256M

解题思路

n^2的DP很容易得出来，在此基础上优化。
b[i]表示前1~i个数中0、1个数差，则f[i]=sigma(f[j]) (|b[j]-b[i]<=k|)。
建一棵线段树，表示b[i]这个位置上f[i]的和，无脑查找即可。

#include
#include
#include
#define maxn 100006
#define fr(i,a,b) for(i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const ll ding=1e9+7;

int i,n,k,a[maxn],b[maxn];
ll s,f[maxn],tr[maxn*8];
void modify(int v,int st,int en,int x,int y)
{
    if (st==en)
    {
        tr[v]=(tr[v]+y)%ding;
        return;
    }
    int m=(st+en) >> 1;
    if (x<=m) modify(v+v,st,m,x,y);
    else modify(v+v+1,m+1,en,x,y);
    tr[v]=(tr[v+v]+tr[v+v+1])%ding;
    return;
}
void findd(int v,int st,int en,int l,int r)
{
    if (st==l && en==r)
    {
        s=(s+tr[v])%ding;
        return;
    }
    int m=(st+en) >> 1;
    if (r<=m) findd(v+v,st,m,l,r);
    else if (l>m) findd(v+v+1,m+1,en,l,r);
    else
    {
        findd(v+v,st,m,l,m);
        findd(v+v+1,m+1,en,m+1,r);
    }
    return;
}
int main()
{
    scanf("%d%d",&n,&k);
    fr(i,1,n) 
    {
        scanf("%d",&a[i]);
        if (a[i]) b[i]=b[i-1]+1;
        else b[i]=b[i-1]-1;
    }
    f[0]=1;
    modify(1,0,n+n,n,1);
    fr(i,1,n)
    {
        int t=b[i]+n;
        s=0;
        findd(1,0,n+n,max(0,t-k),min(t+k,n+n));
        f[i]=s;
        modify(1,0,n+n,t,f[i]);
    }
    printf("%lld\n",f[n]);
    return 0;
}