字符串（后缀数组）：HAOI2016 找相同子串

[HAOI2016]找相同子串

【题目描述】

给定两个字符串，求出在两个字符串中各取出一个子串使得这两个子串相同的方案数。两个方案不同当且仅当这两个子串中有一个位置不同。

【输入格式】

两行，两个字符串s1,s2，长度分别为n1,n2。

【输出格式】

输出一个整数表示答案。

【样例输入】

aabb
bbaa

【样例输出】

10

【数据范围】

对于20%的数据，满足1≤n1,n2≤500；

对于40%的数据，满足1≤n1,n2≤5000；

对于100%的数据，满足1≤n1,n2≤200000，字符串中只有小写字母。

　　和POJ 3415几乎一样。

#include <iostream>
#include <cstring>
#include <cstdio>
using namespace std;
const int N=400010;
char s[N];int len,d;
int r[N],Wa[N],Wb[N],sa[N],rk[N];
int Ws[N],Wv[N],lcp[N];

bool cmp(int *p,int i,int j,int l){
    return p[i]==p[j]&&p[i+l]==p[j+l];
}

void DA(int n,int m){
    int i,j,p,*x=Wa,*y=Wb;
    for(i=0;i<m;i++)Ws[i]=0;
    for(i=0;i<n;i++)++Ws[x[i]=r[i]];
    for(i=1;i<m;i++)Ws[i]+=Ws[i-1];
    for(i=n-1;i>=0;i--)sa[--Ws[x[i]]]=i;
    for(j=1,p=1;p<n;j<<=1,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++)Ws[i]=0;
        for(i=0;i<n;i++)++Ws[Wv[i]=x[y[i]]];
        for(i=1;i<m;i++)Ws[i]+=Ws[i-1];
        for(i=n-1;i>=0;i--)sa[--Ws[Wv[i]]]=y[i];
        for(swap(x,y),p=1,x[sa[0]]=0,i=1;i<n;i++)
            x[sa[i]]=cmp(y,sa[i-1],sa[i],j)?p-1:p++;
    }
}

void LCP(int n){
    int i,j,k=0;
    for(i=1;i<=n;i++)rk[sa[i]]=i;
    for(i=0;i<n;lcp[rk[i++]]=k)
        for(k?k--:k,j=sa[rk[i]-1];r[i+k]==r[j+k];k++);
}

int ID(int x){x=sa[x];
    if(x==d)return 0;
    if(x<d)return 1;
    return 2; 
}

int st[N]={-1},top;
long long sum[N],tot[N],ans;
int main(){
    freopen("find_2016.in","r",stdin);
    freopen("find_2016.out","w",stdout);
    scanf("%s",s);
    len=d=strlen(s);
    scanf("%s",s+len+1);
    s[d]='#';len=strlen(s);
    for(int i=0;i<len;i++)r[i]=s[i];
    DA(len+1,128);LCP(len);
    for(int i=1;i<=len;i++){
        if(!ID(i))continue;
        if(ID(i)==2)ans+=tot[top];
        if(i!=len){
            st[++top]=lcp[i+1];sum[top]=2-ID(i);
            tot[top]=st[top]*sum[top]+tot[top-1];
            while(st[top]<=st[top-1]){
                sum[top-1]+=sum[top];
                st[top-1]=st[top];top--;
                tot[top]=st[top]*sum[top]+tot[top-1];
            }
        }
    }
    top=0;
    for(int i=1;i<=len;i++){
        if(!ID(i))continue;
        if(ID(i)==1)ans+=tot[top];
        if(i!=len){
            st[++top]=lcp[i+1];sum[top]=ID(i)-1;
            tot[top]=st[top]*sum[top]+tot[top-1];
            while(st[top]<=st[top-1]){
                sum[top-1]+=sum[top];
                st[top-1]=st[top];top--;
                tot[top]=st[top]*sum[top]+tot[top-1];
            }
        }
    }
    printf("%lld
",ans);
    return 0;
}