bzoj 3998

后缀自动机上dfs，查询第k大子串

注意代码复杂度，可以不需要加边，写简洁的代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define maxn 1000020

struct node{
    int val,pnt,size,degree;
    int next[27];
}sam[maxn];
struct node2{
    int to,next;
}e[maxn * 2];
int tot,n,k,cnt,t,last,root,head[maxn];
long long sum[maxn];
int pnt[maxn];
char data[500020];
int q[maxn],hh,tt,size[maxn];

inline void add(int x){
    int np = ++tot;
    int p = last;
    sam[np].val = sam[p].val + 1;
    while ( p && !sam[p].next[x] ){
        sam[p].next[x] = np , p = sam[p].pnt;
        sam[np].degree++;
    }
    int q = sam[p].next[x];
    if ( !q ){
        sam[p].next[x] = np;
        sam[np].pnt = p;
        sam[np].degree++;
    }
    else if ( q && sam[q].val == sam[p].val + 1 ){
        sam[np].pnt = q;
    }    
    else{
        int nq = ++tot;
        sam[nq].pnt = sam[q].pnt;
        sam[np].pnt = sam[q].pnt = nq;
        sam[nq].val = sam[p].val + 1;
        //memcpy(sam[nq].next,sam[q].next,sizeof(sam[q].next));
        for (int i = 0 ; i < 26 ; i++){
            if ( sam[q].next[i] ){
                sam[nq].next[i] = sam[q].next[i];
                sam[sam[q].next[i]].degree++;
            }
        }
        while ( p && sam[p].next[x] == q ){
            sam[p].next[x] = nq , p = sam[p].pnt;
            sam[nq].degree++;
            sam[q].degree--;
        }
        if ( sam[p].next[x] == q ){
            sam[p].next[x] = nq;
            sam[nq].degree++ , sam[q].degree--;
        }
    }
    last = np;
}
inline void addedge(int x,int y){
    e[++cnt].to = y;
    e[cnt].next = head[x];
    head[x] = cnt;
}
inline void bfs(int now){
    tt = hh = 0;
    q[tt++] = now;
    while ( hh < tt ){
        now = q[hh++];
        for (int i = head[now] ; i ; i = e[i].next){
            if ( pnt[now] == e[i].to ) continue;
            q[tt++] = e[i].to;
        }
    }
    for (int i = tt - 1 ; i >= 0; i--){
        size[q[i]] += 1;
        size[pnt[q[i]]]    += size[q[i]];
    }
}
void topo(){
    tt = hh = 0;
    for (int i = 0 ; i <= tot ; i++){
        if ( !sam[i].degree ) q[tt++] = i;
    }
    while ( hh < tt ){
        int now = q[hh++];
        for (int i = 0 ; i < 26 ; i++){
            if ( sam[now].next[i] ){
                int to = sam[now].next[i];
                sam[to].degree--;
                if ( sam[to].degree == 0 ) q[tt++] = to;
            }
        }
    }
}
inline void getsum1(){
    for (int i = tt - 1 ; i >= 0 ; i--){
        sum[q[i]] = 1;
        for (int j = 0 ; j < 26 ; j++){
            if ( sam[q[i]].next[j] ){
                sum[q[i]] += sum[sam[q[i]].next[j]];
            }
        }
    }
}
inline void getsum2(){
    for (int i = tt - 1 ; i >= 0 ; i--){
        sum[q[i]] = size[q[i]];
        for (int j = 0 ; j < 26 ; j++){
            if ( sam[q[i]].next[j] ){
                sum[q[i]] += sum[sam[q[i]].next[j]];
            }
        }
    }
}
inline void dfs1(int now,long long k){
    k--;
    if ( k <= 0 ) return;
    for (int i = 0 ; i < 26 ; i++){
        if ( sam[now].next[i] ){
            if ( sum[sam[now].next[i]] >= k ){
                printf("%c",char(i + 'a'));
                   dfs1(sam[now].next[i],k);
                break;
            }
            else k -= sum[sam[now].next[i]];
        }
    }
}
inline void dfs2(int now,long long k){
    k -= (long long) size[now];
    if ( k <= 0 ) return;
    for (int i = 0 ; i < 26 ; i++){
        if ( sam[now].next[i] ){
            if ( sum[sam[now].next[i]] >= k ){
                printf("%c",char(i + 'a'));
                   dfs2(sam[now].next[i],k);
                break;
            }
            else k -= sum[sam[now].next[i]];
        }
    }
}
int main(){
//    freopen("input.txt","r",stdin);
    int __size__ = 50 << 20; // 25MB
    char *__p__ = (char*)malloc(__size__) + __size__;
    __asm__("movl %0, %%esp
" :: "r"(__p__));    
    scanf("%s",data + 1);
    n = strlen(data + 1);
    scanf("%d %d",&t,&k);
    for (int i = 1 ; i <= n ; i++){
        add(int(data[i] - 'a'));
    }
    for (int i = 1 ; i <= tot ; i++){
        pnt[i] = sam[i].pnt;
        addedge(pnt[i],i);
        addedge(i,pnt[i]);
    }
    bfs(root);
    topo();
//    for (int i = 0 ; i <= tot ; i++) cout<<size[i]<<" ";
//    cout<<endl;
    if ( t == 0 ){
        getsum1();
        long long bl = 0;
        for (int i = 1 ; i <= tot ; i++){
            bl += sam[i].val - sam[sam[i].pnt].val;
        }
        if ( bl >= k ) dfs1(root,(long long) k + 1);
        else printf("-1
");
    }
    else{
        getsum2();
        long long bl = (long long) n * ((long long) n - 1) / 2;
        if ( bl >= k ) dfs2(root,(long long) k + size[root]);
        else printf("-1
");
    }
    return 0;
}