HDU - 6194 string string string(后缀数组+RMQ+容斥)

题目链接：点击查看

题目大意：给出一个字符串和一个数字 k，问字符串中出现次数恰好等于 k 次的字串有多少个

题目分析：在跑完后缀数组后，我们可以用sa数组求解，具体做法是枚举起点，找长度为 k 的sa数组的最长前缀就是出现至少k次的答案了，但是不能保证所有子串恰好出现了k次，所以我们在累加答案后再减去出现至少k+1次的答案，最后因为容斥原理还需要加上一部分，因为出现k+1次，多出来的一次可以是向后扩展一个单位，也可以是向前扩展一个单位，注意特判k==1的情况，因为不能在l==r时求lcp，但此时的lcp显然是 len - sa[l] ，特判一下就好了

代码：

#include<iostream>
#include<cstdio> 
#include<string>
#include<ctime>
#include<cstring>
#include<algorithm>
#include<stack>
#include<queue>
#include<map>
#include<set>
#include<cmath>
#include<sstream>
using namespace std;
 
typedef long long LL;
 
const int inf=0x3f3f3f3f;
 
const int N=1e5+100;
 
char str[N];
 
int sa[N]; //SA数组，表示将S的n个后缀从小到大排序后把排好序的
//的后缀的开头位置顺次放入SA中
int t1[N],t2[N],c[N];
 
int rk[N],height[N],len;
 
int s[N];

int st[N][20];
 
void build_sa(int s[],int n,int m)//n为添加0后的总长
{
    int i,j,p,*x=t1,*y=t2;
    for(i=0;i<m;i++) 
        c[i]=0;
    for(i=0;i<n;i++) 
        c[x[i]=s[i]]++;
    for(i=1;i<m;i++) 
        c[i]+=c[i-1];
    for(i=n-1;i>=0;i--) 
        sa[--c[x[i]]]=i;
    for(j=1;j<=n;j<<=1) 
    {
        p=0;
        for(i=n-j;i<n;i++) 
            y[p++]=i;
        for(i=0;i<n;i++) 
            if(sa[i]>=j) 
                y[p++]=sa[i]-j;
        for(i=0;i<m;i++) 
            c[i]=0;
        for(i=0;i<n;i++) 
            c[x[y[i]]]++;
        for(i=1;i<m;i++) 
            c[i]+=c[i-1];
        for(i=n-1;i>=0;i--) 
            sa[--c[x[y[i]]]]=y[i];
        swap(x,y);
        p=1,x[sa[0]]=0;
        for(i=1;i<n;i++) 
            x[sa[i]]=y[sa[i-1]]==y[sa[i]]&&y[sa[i-1]+j]==y[sa[i]+j]?p-1:p++;
        if(p>=n) 
            break;
        m=p;
    }
}
 
void get_height(int s[],int n)//n为添加0后的总长
{
    int i,j,k=0;
    for(i=0;i<=n;i++)
        rk[sa[i]]=i;
    for(i=0;i<n;i++) 
    {
        if(k) 
            k--;
        j=sa[rk[i]-1];
        while(s[i+k]==s[j+k]) 
            k++;
        height[rk[i]]=k;
    }
}
 
void solve(int base=128)
{
    build_sa(s,len+1,base);
    get_height(s,len);
}


void ST_build()
{
    for(int i=1;i<=len;i++)
        st[i][0]=height[i];
    for(int i=1;i<=log2(len);i++)
        for(int j=1;j+(1<<i)-1<=len;j++)
            st[j][i]=min(st[j][i-1],st[j+(1<<i-1)][i-1]);
}
 
int ST_query(int a,int b)
{
    int l=rk[a],r=rk[b];
    if(l>r)
        swap(l,r);
    l++;
    int k=log2(r-l+1);
    return min(st[l][k],st[r-(1<<k)+1][k]);
}

int query(int l,int r)
{
	if(l==r)
		return len-sa[r];
	return ST_query(sa[l],sa[r]);
}

int main()
{
//    freopen("input.txt","r",stdin);
//    ios::sync_with_stdio(false);
    int w;
    cin>>w;
    while(w--)
    {
        int k;
        scanf("%d%s",&k,str);
        len=strlen(str);
        for(int i=0;i<len;i++)
            s[i]=str[i]-'a'+1;
        s[len]=0;
        solve(30);
        ST_build();
        LL ans=0;
        for(int i=1;i+k-1<=len;i++)//枚举sa数组区间
        {
            int l=i,r=i+k-1;
            ans+=query(l,r);
            if(l>1)
                ans-=query(l-1,r);
            if(r<len)
                ans-=query(l,r+1);
            if(l>1&&r<len)
                ans+=query(l-1,r+1);
        }
        printf("%lld\n",ans);
    }
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    return 0;
}