【后缀自动机】 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;
}