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

其实两问都差不多。。。
我们令cnt[i]$cnt[i]$为i$i$这个状态在原串中出现的次数。再在每个状态维护一个siz$siz$。若T=0$T=0$，siz[i[$siz[i[$表示i$i$的所有儿子的siz$siz$之和+1$+1$（子树大小），否则siz[i[$siz[i[$表示i$i$的所有儿子的siz$siz$之和+cnt[i]$+cnt[i]$（类似子树大小的东西）。接下来我们用siz$siz$像平衡树一样查询就好了。
虽然思路简单，但是却因为一个sb错误调了快3个小时。。。后缀自动机里面居然有环，直接dfs会爆炸，得开一个vis数组！我可能是写了假的后缀自动机。。。发现其他大佬都是按照长度从长到短更新的，只有我这个蒟蒻写了dfs。。。

#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
const int N=500005;
int n,last,tot,t,k,fail[N*2],len[N*2],cnt[N*2],ch[N*2][26];
int a[N*2],c[N*2];
ll siz[N*2],sum[N*2];
char s[N],str[N];
bool flag,vis[N*2];
void insert(int x){
    int p=last,np=++tot;
    len[np]=len[p]+1;
    last=np;
    cnt[np]=1;
    for(;p&&!ch[p][x];p=fail[p]){
        ch[p][x]=np;
    }
    if(!p){
        fail[np]=1;
    }else{
        int q=ch[p][x];
        if(len[q]==len[p]+1){
            fail[np]=q;
        }else{
            int nq=++tot;
            len[nq]=len[p]+1;
            memcpy(ch[nq],ch[q],sizeof(ch[q]));
            fail[nq]=fail[q];
            fail[q]=fail[np]=nq;
            for(;p&&ch[p][x]==q;p=fail[p]){
                ch[p][x]=nq;
            }
        }
    }
}
void dfs(int u){
    if(vis[u]){
        return;
    }
    vis[u]=true;
    siz[u]=1;
    sum[u]=cnt[u];
    for(int i=0;i<26;i++){
        if(ch[u][i]){
            dfs(ch[u][i]);
            siz[u]+=siz[ch[u][i]];
            sum[u]+=sum[ch[u][i]];
        }
    }
}
void query1(int u,int dep){
    for(int i=0;i<26;i++){
        if(ch[u][i]){
            if(k==1){
                str[dep]='a'+i;
                flag=true;
                break;
            }else if(k<=siz[ch[u][i]]){
                k--;
                str[dep]='a'+i;
                query1(ch[u][i],dep+1);
                break;
            }else{
                k-=siz[ch[u][i]];
            }
        }
    }
}
void query2(int u,int dep){
    for(int i=0;i<26;i++){
        if(ch[u][i]){
            if(k<=cnt[ch[u][i]]){
                str[dep]='a'+i;
                flag=true;
                break;
            }else if(k<=sum[ch[u][i]]){
                k-=cnt[ch[u][i]];
                str[dep]='a'+i;
                query2(ch[u][i],dep+1);
                break;
            }else{
                k-=sum[ch[u][i]];
            }
        }
    }
}
int main(){
    scanf("%s%d%d",s,&t,&k);
    n=strlen(s);
    last=tot=1;
    for(int i=0;i<n;i++){
        insert(s[i]-'a');
    }
    for(int i=1;i<=tot;i++){
        c[len[i]]++;
    }
    for(int i=1;i<=tot;i++){
        c[i]+=c[i-1];
    }
    for(int i=1;i<=tot;i++){
        a[c[len[i]]--]=i;
    }
    for(int i=tot;i>=1;i--){
        cnt[fail[a[i]]]+=cnt[a[i]];
    }
    dfs(1);
    if(!t){
        query1(1,0);
    }else{
        query2(1,0);
    }
    if(flag){
        puts(str);
    }else{
        puts("-1");
    }
    return 0;
}