poj 3415 Common Substrings

题目链接：http://poj.org/problem?id=3415

题目分类：后缀数组

题意：给出两个串和一个数字k，求两个串的公共字串大于等于k的数目

代码：

//#include<bits/stdc++.h>
#include<stdio.h>
#include<math.h>
#include<algorithm>
#include<string.h>

using namespace std;

#define N 200005
#define LL long long

int wa[N],wb[N],wm[N],wv[N],sa[N];
int *rank,height[N],s[N],a[N];
//sa:字典序中排第i位的起始位置在str中第sa[i]
//rank:就是str第i个位置的后缀是在字典序排第几
//height：字典序排i和i-1的后缀的最长公共前缀

bool cmp(int *r,int a,int b,int l)
{
    return r[a] == r[b] && r[a+l] == r[b+l];
}

void getsa(int *r,int *sa,int n,int m)
{
    int *x=wa,*y=wb,*t;
    for(int i=0; i<m; ++i)wm[i]=0;
    for(int i=0; i<n; ++i)wm[x[i]=r[i]]++;
    for(int i=1; i<m; ++i)wm[i]+=wm[i-1];
    for(int i=n-1; i>=0; --i)sa[--wm[x[i]]]=i;
    for(int i=0,j=1,p=0; p<n; j=j*2,m=p)
    {
        for(p=0,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)wm[i]=0;
        for(i=0; i<n; ++i)wm[x[y[i]]]++;
        for(i=1; i<m; ++i)wm[i]+=wm[i-1];
        for(i=n-1; i>=0; --i)sa[--wm[x[y[i]]]]=y[i];
        for(t=x,x=y,y=t,i=p=1,x[sa[0]]=0; i<n; ++i)
        {
            x[sa[i]]=cmp(y,sa[i],sa[i-1],j)?p-1:p++;
        }
    }
    rank=x;
}

void getheight(int *r,int *sa,int n)
{
    for(int i=0,j=0,k=0; i<n; height[rank[i++]]=k)
    {
        for(k?--k:0,j=sa[rank[i]-1]; r[i+k] == r[j+k]; ++k);
    }
}
int k;
char s1[N];
int len1;

LL solve(int n,int len,int k)
{
    int *mark=wa,*sta=wb,top=0,i;
    LL sum=0,num[3]= {0};
    for(i = 1;i<=n;i++)
    {
        if(height[i]<k)
        {
            top = num[1] = num[2] =0;
        }
        else
        {
            for(int size = top; size&&sta[size]>height[i]-k+1; size--)
            {
                num[mark[size]] += height[i]-k+1-sta[size];
                sta[size] = height[i]-k+1;
            }
            sta[++top] = height[i]-k+1;
            if(sa[i-1]<len) mark[top] = 1;
            if(sa[i-1]>len) mark[top] = 2;
            num[mark[top]]+=height[i]-k+1;
            if(sa[i]<len) sum+=num[2];
            if(sa[i]>len) sum+=num[1];
        }
    }
    return sum;
}

int main()
{
    int i,j;
    while(~scanf("%d",&k),k)
    {
        scanf("%s",s1);
        int n = 0;
        for(n = 0;s1[n]!='\0';n++)
            s[n] = s1[n];
        s[len1=n] = '#';
        scanf("%s",s1+n+1);
        n++;
        for(;s1[n]!='\0';n++)
            s[n] = s1[n];
        s[n] = 0;
        getsa(s,sa,n+1,201);
        getheight(s,sa,n);
        printf("%lld\n",solve(n,len1,k));
    }
    return 0;
}