[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(n2) 做法就是枚举上述的 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;
}