[BZOJ3998][TJOI2015]弦论

[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

Solution
T=0把每个状态的right集合大小视为1，否则就在parent树上递推一遍得到真正的right集合大小
然后在trans图上递推出每个状态往后转移能达到多少子串即可

Code

#include <bits/stdc++.h>
using namespace std;
#define right Right

const int MaxN=1000100;

char s[MaxN];
int right[MaxN],len[MaxN],nxt[MaxN][26],fa[MaxN],f[MaxN];
int w[MaxN],q[MaxN];
int n,T,K;

struct sam{
    int root,last,cnt;
    sam(){root=last=++cnt;}
    void insert(int c){
        int np=++cnt,p=last;last=np;
        right[np]=1;len[np]=len[p]+1;
        for(;p&&!nxt[p][c];nxt[p][c]=np,p=fa[p]);
        if (!p) fa[np]=root;
        else if (len[nxt[p][c]]==len[p]+1) fa[np]=nxt[p][c];
        else{
            int nq=++cnt,q=nxt[p][c];len[nq]=len[p]+1;
            memcpy(nxt[nq],nxt[q],sizeof nxt[q]);fa[nq]=fa[q];
            fa[np]=fa[q]=nq;
            for(;p&&nxt[p][c]==q;nxt[p][c]=nq,p=fa[p]);
        }
    }
    void build(){
        scanf("%s",s+1);n=strlen(s+1);
        for(int i=1;i<=n;i++)insert(s[i]-'a');                
    }
    void topsort(){
        for(int i=1;i<=cnt;i++)w[len[i]]++;
        for(int i=1;i<=n;i++)w[i]+=w[i-1];
        for(int i=1;i<=cnt;i++)q[w[len[i]]--]=i;
    }
    void pre(){
        if (T==0){for(int i=1;i<=cnt;i++)right[i]=1;}
        else{for(int i=cnt,u=q[i];i>=2;i--,u=q[i])right[fa[u]]+=right[u];}
        for(int i=cnt,u=q[i];i>=1;i--,u=q[i]){
            f[u]+=right[u];
            for(int ch=0;ch<26;ch++) if(nxt[u][ch]) f[u]+=f[nxt[u][ch]];
        }
    }
    void cal(){
        for(int p=root;K;){
            for(int ch=0;ch<26;ch++) if (nxt[p][ch]){
                if (f[nxt[p][ch]]>=K){
                    putchar('a'+ch);K-=right[nxt[p][ch]];p=nxt[p][ch];
                    break;
                }else K-=f[nxt[p][ch]];
            }
        }    
    }
}SAM;
int main(){
    SAM.build();
    scanf("%d%d",&T,&K);
    SAM.topsort();SAM.pre();SAM.cal();puts("");
    return 0;
}