POJ 3415 Common Substrings

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

cnt 代表节点所接受的字符在原串中出现的次数
ml 表示该状态能接受的最长的字符串长度
那么我们可以知道一个状态(设为rt)所能接受的最短的字符串长度为ml[pa[rt]]+1
设tmp为当前状态所能接受的字符串长度
那么,一个节点rt，所能接受的字符串在原字符串中所能接受的次数就是cnt[rt]*(tmp-ml[pa[rt]]+1)
那么，因为我们限制了匹配长度k,那么我们要得到max(k,ml[pa[rt]]+1)就是我们要匹配的长度

首先，1 我们要把第一个串建立好SAM，之后拓扑排序一下，记录一下每个状态出现的次数cnt，之后把第二个串在第一个串的SAM上跑一遍，如果匹配长度大于k了，我们要进行一波计数，但是有个坑点，就是我们当前节点大于k但如果我们当前节点的父亲节点的最大长度也大于k了，我们也要记录这个情况，有点类似于线段树的lazy标记。

之后因为我们有了拓扑序了，之后我们按照拓扑序把lazy标记的结果记录一下并且 下放。

#include <iostream>
#include <stdio.h>
#include <string.h>
#define maxs 202020
#define mme(i,j) memset(i,j,sizeof(i))
#define ll long long int
using namespace std;
char a[maxs],b[maxs];
int last,tot,p,q,np,nq;
int son[maxs][53],ml[maxs],pa[maxs];


int newnode(int x)
{
    ml[++tot]=x;
    return tot;
}

void add(int x)
{
    p=last;
    np=newnode(ml[p]+1);

    for(;p&&!son[p][x];p=pa[p] )
        son[p][x]=np;

    if(!p)
        pa[np]=1;
    else {

            q=son[p][x];
        if(ml[q]==ml[p]+1) pa[np]=q;
        else {

            nq=newnode( ml[p]+1 );
            memcpy( son[nq], son[q], sizeof(son[nq]) );
            pa[nq] = pa[q];
            pa[q]  = pa[np]=nq;

            for(;p&&son[p][x]==q;p=pa[p])
                son[p][x]=nq;
        }

    }
    last=np;
}


int tp[maxs],atmp[maxs];


void topo()
{
    for(int i=0;i<=tot;i++) atmp[i]=0;
    for(int i=1;i<=tot;i++) atmp[ml[i]]++;
    for(int i=1;i<=tot;i++) atmp[i]+=atmp[i-1];
    for(int i=1;i<=tot;i++) tp[atmp[ml[i]]--]=i;
}

int cnt[maxs];

void Cntnum()
{
    mme(cnt,0);
    int x,rt=1;

    for(int i=0;a[i];i++)
    {
        if(a[i]>='a'&&a[i]<='z')
                x=a[i]-'a';
        else x=a[i]-'A'+26;

       // printf("i is %d->%c  rt is %d---->son[rt][x] is %d\n",i,a[i],rt,son[rt][x]);

        cnt[son[rt][x]]++;
        rt=son[rt][x];
    }

    for(int i=tot;i>=1;i--)
    {

        rt=tp[i];
        cnt[pa[rt]]+=cnt[rt];
    }

}


int lazy[maxs];

void Query(int k)
{
    mme(lazy,0);
    __int64 ans=0,tmp=0;
    int rt=1,x;

    for(int i=0;b[i];i++)
    {
        if(b[i]>='a'&&b[i]<='z')
                x=b[i]-'a';
        else x=b[i]-'A'+26;

        if( son[rt][x] ){
            tmp++;
            rt = son[rt][x];
        }else {//失配

            while(rt&&!son[rt][x])
                rt=pa[rt];

            if(rt==0) rt=1,tmp=0;
            else tmp = ml[rt]+1,rt=son[rt][x];
        }

        if(tmp>=k)
        {
            ans+= (ll)( tmp - max( k , ml[ pa[rt] ] + 1 ) + 1 ) * cnt[rt];
            if( ml[ pa[rt] ] >= k )
                 lazy[ pa[rt] ]++;
        }
    }

    for(int i=tot;i>=1;i--)
    {
        rt  = tp[i];
        ans+= (ll)lazy[rt]*(ml[rt] - max( k , ml[ pa[rt] ] + 1 ) + 1 )*cnt[rt];

        if(ml[pa[rt]]>=k)
            lazy[pa[rt]]+=lazy[rt];
    }
    printf("%I64d\n",ans);
}

void init()
{
    last=tot=1;
    mme(son,0);
    mme(pa,0);
    //mme(cnt,0);
}

int main()
{
    int k;
    while(~scanf("%d",&k)&&k)
    {
        int x;
        init();
        scanf("%s%s",a,b);
        for(int i=0;a[i];i++)
        {
            if(a[i]>='a'&&a[i]<='z')
                x=a[i]-'a';
            else x=a[i]-'A'+26;
            add(x);
        }
        topo();
        Cntnum();
        Query(k);
    }
    return 0;
}

/*
1
a
a
1
aa
aa
1
aaa
aaa
1
aab
baa
*/