BZOJ_3998_[TJOI2015]弦论_后缀自动机

#include <cstdio>
#include <string.h>
#include <algorithm>
using namespace std;
#define N 1000050
int ch[N][26],fa[N],dep[N],siz[N],cnt=1,lst=1,f[N];
int ws[N],a[N];
char w[N];
void insert(int x) {
	int p=lst,np=++cnt,q,nq;
	lst=np; dep[np]=dep[p]+1;
	for(;p&&!ch[p][x];p=fa[p]) ch[p][x]=np;
	if(!p) fa[np]=1;
	else {
		q=ch[p][x];
		if(dep[q]==dep[p]+1) fa[np]=q;
		else {
			fa[nq=++cnt]=fa[q];
			dep[nq]=dep[p]+1;
			memcpy(ch[nq],ch[q],sizeof(ch[q]));
			fa[q]=fa[np]=nq;
			for(;p&&ch[p][x]==q;p=fa[p]) ch[p][x]=nq;
		}
	}
	siz[np]=1;
}
void print() {
	int i,j;
	printf("test-------------------------------------------\n");
	for(i=1;i<=cnt;i++) {
		printf("p=%d,siz=%d,dep=%d,fa=%d\n",i,siz[i],dep[i],fa[i]);
		for(j=0;j<26;j++) {
			if(ch[i][j]) {
				printf("ch(%d)(%c)=%d\n",i,j+‘a‘,ch[i][j]);
			}
		}
	}
	printf("lst=%d\n",lst);
}
int main() {
	int T,K;
	scanf("%s%d%d",w+1,&T,&K);
	int n=strlen(w+1);
	int i,j,p;
	for(i=1;i<=n;i++) insert(w[i]-‘a‘);
	for(i=1;i<=cnt;i++) ws[dep[i]]++;
	for(i=1;i<=n;i++) ws[i]+=ws[i-1];
	for(i=cnt;i;i--) a[ws[dep[i]]--]=i;
	if(T) {
		for(i=cnt;i;i--) {
			p=a[i]; siz[fa[p]]+=siz[p];
		}
	}else for(i=cnt;i;i--) siz[a[i]]=1;
	for(siz[1]=0,i=cnt;i;i--) {
		p=a[i];
		for(f[p]=siz[p],j=0;j<26;j++) f[p]+=f[ch[p][j]];
	}
	if(f[1]<K) puts("-1");
	else {
		p=1;
		while(K>siz[p]) {
			K-=siz[p];
			for(i=0;i<26;i++) {
				if(K<=f[ch[p][i]]) {printf("%c",i+‘a‘),p=ch[p][i]; break;}
				else K-=f[ch[p][i]];
			}
		}
		puts("");
	}
}