POJ 3415 Common Substrings (求长度不小于k的公共子串的个数)

/*
 * POJ 3415 Common Substrings
 * 给出两个字符串A，B，求长度不小于k的公共子串的个数
 *
 * 将B接到A的后面，中间用一个没出现的字符隔开，然后求sa和height数组
 * 然后按height分组，我们只考虑height[i]>=k的后缀，因为小于不符合题意，然后对于每个分组，统计答案
 * 很容易想到算法就是对于同一组里面的B后缀，我和前面的所有A的后缀比较，累加答案，再反过来对A做一遍即可
 * 这样的时间复杂度是O(n*n)的，会TLE
 * 考虑这样的一个性质，后缀i和j的lcp是height[rank[i+1]]...height[rank[j]]的最小值，那么对于
 * height[i]==s&&sa[i]是B的后缀，它和前面A的答案肯定是最小的height的累加，于是这样我们可以用一个
 * 单调栈来维护，每当新加入的height小于栈顶的元素，我们就把栈顶元素弹出来，把新的值压入栈中
 * 怎么统计答案呢？可以先加后减，当计算B的个数时，遇到A的后缀我们就直接累加到答案中，如果和栈顶元素冲突了
 * 说明多加了一部分，在弹出元素时减去即可，这样多维护一个个数就行了。
 */

#include <stdio.h>
#include <string.h>
#include <iostream>
using namespace std;

const int MAXN = 200000+100;

int sa[MAXN];
int t1[MAXN],t2[MAXN],c[MAXN];
int Rank[MAXN],height[MAXN];
void build_sa(int s[],int n,int m)
{
    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 getHeight(int s[],int n)
{
    int i,j,k=0;
    for(i=0;i<=n;i++) Rank[sa[i]]=i;
    for(i=0;i<n;i++)
    {
        if(k)k--;
        j=sa[Rank[i]-1];
        while(s[i+k]==s[j+k])k++;
        height[Rank[i]]=k;
    }
}
char str1[MAXN],str2[MAXN];
int ss[MAXN];
int sta[MAXN],stb[MAXN];

int main()
{
    int k;
    while(scanf("%d",&k)==1&&k)
    {
        scanf("%s",str1);
        scanf("%s",str2);
        int len1=strlen(str1);
        int len2=strlen(str2);
        int n=len1+len2+1;
        for(int i=0;i<len1;i++) ss[i]=str1[i];
        ss[len1]=1;
        for(int i=0;i<len2;i++) ss[len1+i+1]=str2[i];
        ss[n]=0;
        build_sa(ss,n+1,128);
        getHeight(ss,n);
        long long ans=0;
        long long res=0;
        int top=0;
        for(int i=2;i<=n;i++)
        {
            if(height[i]<k)
            {
                res=0;
                top=0;
                continue;
            }
            int cnt=0;
            if(sa[i-1]<len1)
            {
                cnt++;
                res+=height[i]-k+1;
            }
            while(top>0&&height[i]<=sta[top-1])
            {
                top--;
                res-=stb[top]*(sta[top]-height[i]);
                cnt+=stb[top];
            }
            sta[top]=height[i],stb[top++]=cnt;
            if(sa[i]>len1) ans+=res;
        }
        res=0,top=0;
        for(int i=2;i<=n;i++)
        {
            if(height[i]<k)
            {
                res=0;
                top=0;
                continue;
            }
            int cnt=0;
            if(sa[i-1]>len1)
            {
                cnt++;
                res+=height[i]-k+1;
            }
            while(top>0&&height[i]<=sta[top-1])
            {
                top--;
                res-=stb[top]*(sta[top]-height[i]);
                cnt+=stb[top];
            }
            sta[top]=height[i],stb[top++]=cnt;
            if(sa[i]<len1) ans+=res;
        }
        printf("%lld\n",ans);
    }
    return 0;
}