Codeforces 271D. Good Substrings (后缀数组+前缀和)
题目: http://codeforces.com/problemset/problem/271/D
题意:
给定一个字母串s,给定每个字母是good/bad字母;
一个子串如果包含的bad字母数<=k,则为一个good串;
求有多少个good串。
分析:
首先可以用前缀和处理出s[0]~s[i]有多少个bad字母,用sum[i]表示;
则可以 O(1) 地查询任意一个串是否是good串;
如果枚举起点(L)和终点(R)的话,会统计到重复的子串,比如abceabc会把abc统计两次;
那怎么办呢?
用 后缀数组!
把字符串s用后缀数组处理后,再求得height数组;
其中height[i]表示排名为i的字符串和排名为i-1的字符串的最长公共前缀长度;
比如排名为i-1的串为abcde,排名为i的串为abcdkmn,则height[i]=4;
那么怎么避免重复统计子串呢?
规定子串的左端点为L,右端点为R;
循环L,让其依次等于rank1串的左端点、rank2串的左端点、…、rankn串的左端点;
R则从L+height[i]开始往n循环,每一个(L,R)用前缀和O(1)判断是否good,是则ans++。
代码:
#include <bits/stdc++.h>
using namespace std;
const int tmax=1505;
int t1[tmax*4],t2[tmax*4],c[tmax*4],rrank[tmax];
char s[tmax],good[30];
int k,len,sa[tmax],height[tmax],sum[tmax],ans,maxk;
bool cmp(int *y,int i,int j)
{
if(y[i]!=y[j]) return y[i]<y[j];
int ri=i+k<=len?y[i+k]:-1;
int rj=j+k<=len?y[j+k]:-1;
return ri<rj;
}
void suffix()
{
int *x=t1,*y=t2,m=256;
for(int i=0;i<m;i++) c[i]=0;
for(int i=0;i<len;i++) c[x[i]=s[i]]++;
for(int i=1;i<m;i++) c[i]+=c[i-1];
for(int i=len-1;i>=0;i--) sa[--c[x[i]]]=i;
for(k=1;k<len;k<<=1)
{
int p=0;
for(int i=len-k;i<len;i++) y[p++]=i;
for(int i=0;i<len;i++)
if(sa[i]>=k) y[p++]=sa[i]-k;
for(int i=0;i<m;i++) c[i]=0;
for(int i=0;i<len;i++) c[x[y[i]]]++;
for(int i=1;i<m;i++) c[i]+=c[i-1];
for(int i=len-1;i>=0;i--) sa[--c[x[y[i]]]]=y[i];
swap(x,y);
m=1;x[sa[0]]=0;
for(int i=1;i<len;i++)
x[sa[i]]=cmp(y,sa[i-1],sa[i])==true? m++:m-1;
if(m>=len) break;
}
return;
}
void get_height()
{
int i,j,h=0;
for(i=0;i<len;i++) rrank[sa[i]]=i;
height[0]=0;
for(i=0;i<len;i++)
{
if(rrank[i]==0) continue;
j=sa[rrank[i]-1];
if(h>0) h--;
while(i+h<len&&j+h<len&&s[i+h]==s[j+h]) h++;
height[rrank[i]]=h;
}
return;
}
void work()
{
int i,j,ll,rr;
for(i=0;i<len;i++)
sum[i]=sum[i-1]+(good[s[i]-'a']=='0'?1:0);
for(i=1;i<len;i++)
{
ll=sa[i];
rr=sa[i]+height[i];
for(j=rr;j<len-1;j++)
if(sum[j]-sum[ll-1]<=maxk) ans++;
}
printf("%d",ans);
return;
}
int main()
{
scanf("%s",s);
len=strlen(s);
s[len++]=0;
scanf("%s",good);
scanf("%d",&maxk);
suffix();
get_height();
work();
return 0;
}