【BZOJ3998】[TJOI2015]弦论 后缀自动机

【BZOJ3998】[TJOI2015]弦论

Description

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

Input

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

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

Output

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

Sample Input

aabc

0 3

Sample Output

aab

HINT

N<=5*10^5

T<2

K<=10^9

题解：先跑后缀自动机，再处理出SAM的拓扑序，处理出在每个节点结束的子串个数，最后扫一遍SAM。如果在当前儿子的子树中结束的子串个数<k，那么k-=个数，再扫下一个儿子。

#include <cstdio>
#include <cstring>
#include <iostream>
using namespace std;
const int maxn=1000010;
int n,T,K,last,tot;
int st[maxn],tb[maxn],pre[maxn],ch[maxn][26],dep[maxn],r[maxn],f[maxn];
char str[maxn];
void add(int x)
{
    int p=last,np=++tot;
    last=np,dep[np]=dep[p]+1;
    for(;p&&!ch[p][x];p=pre[p]) ch[p][x]=np;
    if(!p)  pre[np]=1;
    else
    {
        int q=ch[p][x];
        if(dep[q]==dep[p]+1)    pre[np]=q;
        else
        {
            int nq=++tot;
            memcpy(ch[nq],ch[q],sizeof(ch[q]));
            dep[nq]=dep[p]+1,pre[nq]=pre[q],pre[q]=pre[np]=nq;
            for(;p&&ch[p][x]==q;p=pre[p])   ch[p][x]=nq;
        }
    }
}
int main()
{
    last=tot=1;
    scanf("%s%d%d",str,&T,&K),n=strlen(str);
    int i,p,j;
    for(i=1;i<=n;i++)    add(str[i-1]-'a');
    for(i=p=1;i<=n;i++)  p=ch[p][str[i-1]-'a'],r[p]++;
    for(i=1;i<=tot;i++)  st[dep[i]]++;
    for(i=1;i<=n;i++)    st[i]+=st[i-1];
    for(i=tot;i>=1;i--)  tb[st[dep[i]]--]=i;
    if(T)   for(i=tot;i>=1;i--)  r[pre[tb[i]]]+=r[tb[i]];
    else    for(i=1;i<=tot;i++)  r[i]=1;
    for(i=tot,r[1]=0;i>=1;i--)   for(f[tb[i]]=r[tb[i]],j=0;j<26;j++)  f[tb[i]]+=f[ch[tb[i]][j]];
    if(f[1]<K)
    {
        printf("-1");
        return 0;
    }
    p=1;
    while(K>r[p])
    {
        for(i=0,K-=r[p];i<26&&K>f[ch[p][i]];i++)  K-=f[ch[p][i]];
        printf("%c",'a'+i),p=ch[p][i];
    }
    return 0;
}