[POJ 3415]Common Substrings(后缀数组)

题目链接

http://poj.org/problem?id=3415

题目大意

求字符串$A与B$的长度大于等于$m$的公共子串个数。

思路

看到公共子串一类的题目，很容易想到用后缀数组做。我们首先把$B$串加到$A$串后面，中间用一个字符$分割开，然后求后缀数组
。对于$A$的后缀$i$和$B$的后缀$j$而言，若它们的$LCP>=m$，则它们会给最终的答案贡献$(i-m+1)(j-m+1)$。一个显然的$O(n^2)$做法就是枚举上述的$i,j$，累计答案。但是这样做显然是不能AC的，考虑用单调栈来对新串的$height$数组进行分段。

在单调栈中保存的东西是$height[i]-m+1$，即，我们标记栈里每个元素分别属于$A$串还是$B$串，保证栈里的所有元素都在原串的$height$数组中是连续的，并且保证对任意的$i，height[i]>=m$。并记录两个变量$num[1]、num[2]$，分别代表$num[1]=$字符串$A$和当前入栈的$suffix(i)$的大于等于$m$的公共子串个数，$num[2]=$字符串$B$和当前入栈的$suffix(i)$的大于等于$m$的公共子串个数。

那么我们假如现在加入的$sa[i]$的$height[i]<m$，则不能保证栈里的元素是连续的，要清空栈及相关数组。

假如我们现在加入的$sa[i]$的$height[i]>=m$,我们可以在栈顶那部分中找到一个区间，使得这个区间内所有的$sa[j]，height[j]-m+1>=height[i]-m+1$，这意味着什么呢？意味着这个区间内所有的后缀$sa[j]$与后缀$sa[i]$的LCP变小了，从$height[i-1]$变成了$height[i]$，每个$sa[j]和sa[i]$配对各能产生的公共子串个数也减少了，减少的子串个数是$(height[j]-m+1)-(height[i]-m+1)$。那么就让$i$所属的串记录减少贡献，$num[mark[i]]-=(height[j]-m+1)-(height[i]-m+1)$，这里$mark[i]=1$表示$sa[i]$是属于$A$串的，$mark[i]=2$表示$sa[i]$属于$B$串。

在栈中插入完$sa[i]$后要更新附加信息，并且由于此时的$num[1]、num[2]$代表的是$A$串和$B$串各自和$suffix(sa[i])$构成的长度大于等于$m$的公共子串个数。若$sa[i]$属于$A$串，那么就在总答案中添加$B$串和$suffix(sa[i])$构成的长度大于等于$m$的公共子串个数，即$num[2]$，反之亦然。

代码

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>

#define MAXN 1100000

using namespace std;

typedef long long int LL;

int sa[MAXN],rank[MAXN],height[MAXN];
int wa[MAXN],wb[MAXN],wv[MAXN],cnt[MAXN];

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

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

void calc(int *r,int n)
{
    int i,j,k=0;
    for(i=1;i<=n;i++) rank[sa[i]]=i;
    for(i=0;i<n;height[rank[i++]]=k)
        for(k?k--:0,j=sa[rank[i]-1];r[i+k]==r[j+k];k++); //!!!!!!
}

char a[MAXN],b[MAXN];
int s[MAXN],n,m,stack[MAXN],top=0;
int mark[MAXN]; //mark[i]=1表示stack[i]是属于A串的，mark[i]=2表示stack[i]属于B串

int main()
{
    while(scanf("%d",&m)!=EOF&&m)
    {
        n=0;
        scanf("%s%s",a,b);
        int lena=strlen(a),lenb=strlen(b);
        for(int i=0;i<lena;i++) s[n++]=a[i];
        s[n++]='$';
        for(int i=0;i<lenb;i++) s[n++]=b[i];
        s[n]=0;
        SA(s,n+1,300);
        calc(s,n);
        top=0;
        LL sum=0,num[3]={0,0,0}; //num[1]=字符串
        for(int i=1;i<=n;i++)
        {
            if(height[i]<m) top=num[1]=num[2]=0; //这一段已经结束了，清空栈和记录信息的数组
            else
            {
                for(int p=top;p>0&&stack[p]>height[i]-m+1;p--)
                {
                    num[mark[p]]+=height[i]-m+1-stack[p];
                    stack[p]=height[i]-m+1;
                }
                stack[++top]=height[i]-m+1;
                if(sa[i-1]<lena) mark[top]=1;
                else if(sa[i-1]>lena) mark[top]=2;
                num[mark[top]]+=height[i]-m+1;
                if(sa[i]<lena) sum+=num[2];
                else if(sa[i]>lena) sum+=num[1];
            }
        }
        printf("%lld\n",sum);
    }
    return 0;
}