bzoj3998 [后缀自动机]

题目

对于一个给定长度为N的字符串，求它的第K小子串是什么。

输入

第一行是一个仅由小写英文字母构成的字符串S

第二行为两个整数T和K，T为0则表示不同位置的相同子串算作一个。T=1则表示不同位置的相同子串算作多个。K的意义如题所述。

输出

输出仅一行，为一个数字串，为第K小的子串。如果子串数目不足K个，则输出-1

样例输入

aabc
0 3

样例输出

aab

限制条件

$N\le 5*10^5$
$T< 2$
$K \le 10^9$

分析

对输入的字符串建一个后缀自动机,分类讨论r[]的维护情况.
若T==0,每个相同子串只算一次,所以r[]数组每个元素都为1,r[tmp[i]]=1;
若T==1,每个相同子串可算多次,从parent树上儿子向父亲维护,r[fa[tmp[i]]]+=r[tmp[i]];
需要注意的是,两种情况下的r[root]=0

完整代码

#include<bits/stdc++.h>
#define maxn 500010
#define maxt 1000010
//#define DEBUG
using namespace std;
int n,sign,K;
int size[maxt],sum[maxt],tmp[maxt];
char ch[maxn];

inline int read()
{
    char ch;
    int sign=1,read=0;
    do
        ch=getchar();
    while((ch<'0'||ch>'9')&&ch!='-');
    if(ch=='-') sign=-1,ch=getchar();
    while(ch>='0' && ch<='9')
    {
        read=read*10+ch-'0';
        ch=getchar();
    }
    return read*sign;
}

struct SAM
{
    int tot,root,last;
    int son[maxt][26],maxl[maxt],fa[maxt],r[maxt];

    void init() { tot=root=last=1; }
    int insert(int x) { return maxl[++tot]=x,tot; }

    void add(int pos)
    {
        int x=ch[pos]-'a',np=insert(pos),p=last;
        last=np,r[np]=1;
        for( ; p&&!son[p][x] ; p=fa[p] ) son[p][x]=np;
        if(!p) fa[np]=root;
        else
        {
            int q=son[p][x];
            if(maxl[q]==maxl[p]+1) fa[np]=q;
            else
            {
                int nq=insert(maxl[p]+1);
                memcpy(son[nq],son[q],sizeof(son[q]));
                fa[nq]=fa[q];
                fa[np]=fa[q]=nq;
                for( ; son[p][x]==q ; p=fa[p] ) son[p][x]=nq;
            }
        }
    }

    void Tsort()
    {
        for(int i=1;i<=tot;i++) sum[maxl[i]]++;
        for(int i=1;i<=n;i++) sum[i]+=sum[i-1];
        for(int i=1;i<=tot;i++) tmp[sum[maxl[i]]--]=i;
    #ifdef DEBUG
        for(int i=1;i<=tot;i++) printf("tmp[%d]=%d  maxl[tmp[%d]]=%d\n",i,tmp[i],i,maxl[tmp[i]]);
    #endif
        for(int i=tot;i;i--)
        {
            int x=tmp[i];
            if(sign==1) r[fa[x]]+=r[x];
            else  r[x]=1; 
        }
        r[root]=0;
        for(int i=tot;i;i--)
        {
            int x=tmp[i];
            size[x]=r[x];
            for(int i=0;i<26;++i)
                if(son[x][i]) size[x]+=size[son[x][i]];
        }
    }

    void build()
    {
        init();
        scanf("%s",ch+1);
        n=strlen(ch+1);
        for(int i=1;i<=n;++i) add(i);
    }

    void dfs(int x,int k)
    {
        if(k<=r[x]) return;
        k-=r[x];
        for(int i=0;i<26;++i)
            if(son[x][i])
            {
                if(k<size[son[x][i]])
                {
                    putchar(i+'a');
                    dfs(son[x][i],k);
                    return;
                }
                else k-=size[son[x][i]];
            }
    }

    void work()
    {
    #ifdef DEBUG
        for(int i=tot;i;i--) printf("size[tmp[%d]]=%d\n",i,size[tmp[i]]);
        for(int i=tot;i;i--) printf("r[tmp[%d]]=%d\n",i,r[tmp[i]]);
    #endif
        if(size[root]<K)
        {
            printf("%d",-1);
            return;
        }
        else 
        {
            dfs(root,K);
            return;
        }
    }

} sam ;

int main()
{
    sam.build();
    sign=read(),K=read();
    sam.Tsort();
    sam.work();
    return 0;
}