[BZOJ2555] [TJOI2015]弦論 && 後綴自動機

本來以爲是一個多麼神的題 弦論 多麼高大上 

結果 居然是後綴自動機的統計 水水的

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<iostream>
#include<queue>
#define SF scanf
#define PF printf
#define idx(c) (c-'a')
using namespace std;
typedef long long LL;
const int MAXN = 500000;
int T, K;
char s[MAXN+10];
struct SAM {
    int fa[MAXN*2+10], ch[MAXN*2+10][26], step[MAXN*2+10];
    int val[MAXN*2+10], Q[MAXN*2+10], tmp[MAXN*2+10], cnt[MAXN*2+10];
    int ncnt, last, length;
    SAM () { last = ++ncnt; }
    void extend(int c) {
        int p = last, np = last = ++ncnt;
        step[np] = step[p]+1; val[np] = 1;
        while(!ch[p][c] && p) ch[p][c] = np, p = fa[p];
        if(!p) fa[np] = 1;
        else {
            int q = ch[p][c], nq;
            if(step[q] == step[p]+1) fa[np] = q;
            else {
                nq = ++ncnt;
                step[nq] = step[p]+1;
                memcpy(ch[nq], ch[q], sizeof(ch[nq]));
                fa[nq] = fa[q];
                fa[q] = fa[np] = nq;
                while(ch[p][c] == q) ch[p][c] = nq, p = fa[p];
            }
        }
    }
    void build(char *s) {
        length = strlen(s);
        for(int i = 0; i < length; i++)
            extend(idx(s[i]));
    }
    void sort() {
        for(int i = 1; i <= ncnt; i++) tmp[step[i]]++;
        for(int i = 1; i <= length; i++) tmp[i] += tmp[i-1];
        for(int i = ncnt; i; i--) Q[tmp[step[i]]--] = i;
    }
    void pre() {
        for(int i = ncnt; i; i--) 
            if(T) val[fa[Q[i]]] += val[Q[i]];
            else val[Q[i]] = 1;
        val[1] = 0;
        for(int i = ncnt; i; i--) {
            int u = Q[i];
            cnt[u] = val[u];
            for(int j = 0; j < 26; j++) cnt[u] += cnt[ch[u][j]];
        }
    }
    void dfs(int u, int K) {
        if(K <= val[u]) return ;
        K -= val[u];
        for(int i = 0; i < 26; i++)
            if(ch[u][i]) {
                if(K <= cnt[ch[u][i]]) {
                    putchar(i+'a');
                    dfs(ch[u][i], K);
                    return ;
                }
                K -= cnt[ch[u][i]];
            }
    }
} sam;
int main() {
    SF("%s", s);
    SF("%d%d", &T, &K);
    sam.build(s);
    sam.sort();
    sam.pre();
    if(K > sam.cnt[1]) puts("-1");
    else sam.dfs(1, K);
    return 0;
}