【HDU - 6194】后缀数组+线段树

题意：给出一个只有小写字母的串，求至少出现k次的Substring有多少个。

 

思路：先求出height数组，然后枚举长度为k的段，查询区间的height（ST表或者线段树），他们的LCP和就是答案，但是这样会算出长度为k+1的和长度为k-1的串，所以要减去这一部分

 

#include <bits/stdc++.h>

using namespace std;


#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
char t[100000+10];

const int maxn = 1e5+10;
const int INF = 0x3f3f3f3f;
int Min[maxn << 2];
int wa[maxn],wb[maxn],wv[maxn],wu[maxn],s[maxn];
int cmp(int *r,int a,int b,int l)
{
    return r[a]==r[b] && r[a+l]==r[b+l];
}
int sa[maxn],rk[maxn],height[maxn];
void SA(int *r,int n,int m)
{
    int *x=wa,*y=wb;

    for(int i=0; i<m; ++i)
        wu[i]=0;
    for(int i=0; i<n; ++i)
        ++wu[x[i]=r[i]];
    for(int i=1; i<m; ++i)
        wu[i]+=wu[i-1];
    for(int i=n-1; i>=0; --i)
        sa[--wu[x[i]]]=i;

    int p=1;
    for(int j=1; p<n; j<<=1,m=p)
    {
        p=0;
        for(int i=n-j; i<n; ++i)
            y[p++]=i;
        for(int i=0; i<n; ++i)
            if(sa[i]>=j)
                y[p++]=sa[i]-j;
        for(int i=0; i<n; ++i)
            wv[i]=x[y[i]];
        for(int i=0; i<m; ++i)
            wu[i]=0;
        for(int i=0; i<n; ++i)
            ++wu[wv[i]];
        for(int i=1; i<m; ++i)
            wu[i]+=wu[i-1];
        for(int i=n-1; i>=0; --i)
            sa[--wu[wv[i]]]=y[i];
        swap(x,y);
        x[sa[0]]=0;
        p=1;
        for(int i=1; i<n; ++i)
            x[sa[i]]=cmp(y,sa[i-1],sa[i],j)?p-1:p++;
    }
    for(int i=1; i<n; ++i)
        rk[sa[i]]=i;
    int k=0;
    for(int i=0; i<n-1; height[rk[i++]]=k)
    {
        if(k)
            --k;
        for(int j=sa[rk[i]-1]; r[i+k]==r[j+k]; ++k);
    }
//    cout << "enter" << endl;
}


void pushup(int rt)
{
    Min[rt] = min(Min[rt<<1],Min[rt<<1|1]);
}
void build(int l, int r, int rt)
{
    if(l == r)
    {
        Min[rt] = height[l];
        return;
    }
    int m = (l + r) >> 1;
    build(lson);
    build(rson);
    pushup(rt);
}
int query(int  L, int R, int  l, int r, int rt)
{
    if(L <= l && r <= R)
    {
        return Min[rt];
    }
    int m = (l + r) >> 1;
    int ret = INF;
    if(L<=m)
        ret = min(ret,query(L,R,lson));
    if(R>m)
        ret = min(ret,query(L,R,rson));
    return ret;
}
int main()
{
    int T;
    cin >> T;
    while(T--)
    {
        memset(height,0,sizeof(height));
        memset(s,0,sizeof(s));
        memset(wb,0,sizeof(wb));
        memset(wa,0,sizeof(wa));
        memset(wv,0,sizeof(wv));
        memset(wu,0,sizeof(wu));
        int k;
        cin >> k;
        scanf("%s",t);
        int n = strlen(t);
        int len = n;
        for(int i = 0 ; i < n ; i ++)
        {
            s[i] = t[i] - 'a' + 1;
        }
        s[n] = 0;
        SA(s, n+1, 30);
//        ll ans = 0;
        long long ans = 0;
//        if(k >= 2){
        build(1,n,1);
//        k --;
        for(int i = 1 ; i + k - 1 <= n ; i++)
        {
            if(i != i+k-1)
                ans += query(i+1,i+k-1,1,n,1);
            else
                ans += (len-sa[i]);
            if(i - 1 >= 1 && i+k-1 <= n)
            {
                ans -= query(i-1+1,i+k-1,1,n,1);
            }
            if(i+k <= n && i >= 1)
            {
                ans -= query(i+1,i+k,1,n,1);
            }
            if(i - 1 >= 1 && i + k <= n)
                ans += query(i-1+1,i+k,1,n,1);
        }
        }
//        else{
//            for(int i = 1 ; i < n ; i ++){
//                int len = n - sa[i]
//                if()
//            }
//        }
        cout << ans << endl;
    }
    return 0;
}