【后缀自动机】 BZOJ 3998: [TJOI2015]弦论

建立好后缀自动机后，分别讨论一下是要求重复子串还是不重复子串。。。

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;

const int alpha = 26;
const int maxn = 1000005;

struct node
{
	int len;
       	LL cnt, size;
	node *fa, *ch[alpha];
}pool[maxn], *root, *last;

int C[maxn];
int tot;

node* newnode(int len)
{
	node *p = &pool[tot++];
	p->len = len, p->cnt = 0, p->size = 0;
	p->fa = 0;
	memset(p->ch, 0, sizeof p->ch);
	return p;
}

void init()
{
	tot = 0;
	root = last = newnode(0);
}

node* add(int c)
{
	node *p = last, *np = newnode(p->len + 1);
	last = np;
	for(; p && !p->ch[c]; p = p->fa) p->ch[c] = np;
	if(!p) np->fa = root;
	else {
		node *q = p->ch[c];
		if(q->len == p->len + 1) np->fa = q;
		else {
			node *nq = newnode(q->len + 1);
			*nq = *q;
			nq->len = p->len + 1;
			q->fa = np->fa = nq;
			for(; p && p->ch[c] == q; p = p->fa) p->ch[c] = nq;
		}
	}
	return np;
}

node *tp[maxn];
void Tp()
{
	memset(C, 0, sizeof C);
	for(int i = 0; i < tot; i++) C[pool[i].len]++;
	for(int i = 1; i < tot; i++) C[i] += C[i-1];
	for(int i = 0; i < tot; i++) tp[--C[pool[i].len]] = &pool[i];
}

char s[maxn];
vector<char> ans;
int ok;

void dfs(node *p, LL& k)
{
	if(ok) return;
	if(k <= 0) {
		ok = 1;
		return;
	}

	for(int i = 0; i < alpha && !ok; i++) if(p->ch[i]) {
		if(k > p->ch[i]->size) k -= p->ch[i]->size;
		else {
			k -= p->ch[i]->cnt;
			ans.push_back(i + 'a');
			dfs(p->ch[i], k);
		}
	}
}

void work()
{
	int op;
	LL kk;
	init();
	for(int i = 0; s[i]; i++) add(s[i] - 'a');
	
	scanf("%d%lld", &op, &kk);

	if(op) {
		node *p = root;
		for(int i = 0; s[i]; i++) p = p->ch[s[i] - 'a'], p->cnt++;
	}
	else {
		for(int i = 0; i < tot; i++) pool[i].cnt = 1;
	}

	Tp();
	if(op) {
		for(int i = tot - 1; i > 0; i--) tp[i]->fa->cnt += tp[i]->cnt;
	}
	for(int i = tot - 1; i >= 0; i--) {
		tp[i]->size += tp[i]->cnt;
		for(int j = 0; j < alpha; j++) if(tp[i]->ch[j])
			tp[i]->size += tp[i]->ch[j]->size;
	}
	
	LL res = 0;
	for(int i = 0; i < alpha; i++) if(root->ch[i]) res += root->ch[i]->size;
	
	if(res < kk) {
		printf("-1\n");
		return;
	}

	ok = 0;
	ans.clear();
	dfs(root, kk);

	for(int i = 0; i < ans.size(); i++) printf("%c", ans[i]);
	printf("\n");
}

int main()
{
	while(scanf("%s", s) != EOF) {
		work();
	}
	
	return 0;
}