POJ3415:Common Substrings(后缀数组+单调栈)

A substring of a string T is defined as:

Given two strings A, B and one integer K, we define S, a set of triples (i, j, k):

You are to give the value of |S| for specific A, B and K.

The input file contains several blocks of data. For each block, the first line contains one integer K, followed by two lines containing strings A and B, respectively. The input file is ended by K=0.

1 ≤ |A|, |B| ≤ 10⁵
1 ≤ K ≤ min{|A|, |B|}
Characters of A and B are all Latin letters.

For each case, output an integer |S|.

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stack>
#include <queue>
#include <map>
#include <set>
#include <vector>
#include <math.h>
#include <bitset>
#include <algorithm>
#include <climits>
using namespace std;

#define LS 2*i
#define RS 2*i+1
#define UP(i,x,y) for(i=x;i<=y;i++)
#define DOWN(i,x,y) for(i=x;i>=y;i--)
#define MEM(a,x) memset(a,x,sizeof(a))
#define W(a) while(a)
#define gcd(a,b) __gcd(a,b)
#define LL long long
#define N (2*100000+10)
#define MOD 1000000007
#define INF 0x3f3f3f3f
#define EXP 1e-8
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;
}