[BZOJ3998][TJOI2015]弦论

原题地址

求一个串的严格/非严格K小子串.

后缀自动机经典问题.

AC code:

#include <cstdio>
#include <cstring>
typedef long long ll;
const int K=26;
const int N=1000010;
int  t,k,len,cnt;
char s[N];

struct nod{
    bool ed;
    ll   l,f,g;
    nod  *pr,*ch[K];
}pool[N];

struct SAM{
    nod *root,*last;

    SAM(){
        root=last=&pool[cnt++];
    }

    void extend(int x){
        nod *p=last,*np=&pool[cnt++];
        for(last=np,np->l=p->l+1;p&&(!p->ch[x]);p->ch[x]=np,p=p->pr) ;
        if(!p) np->pr=root;
        else{
            nod *q=p->ch[x],*nq;
            if(q->l==p->l+1) np->pr=q;
            else{
                nq=&pool[cnt++];*nq=*q;nq->l=p->l+1;np->pr=q->pr=nq;
                for(;p&&p->ch[x]==q;p->ch[x]=nq,p=p->pr) ;
            }
        }
    }
    void dfs1(nod *x){
        x->f=1;
        for(int i=0;i<K;i++){
            if(!x->ch[i]) continue;
            if(!x->ch[i]->f) dfs1(x->ch[i]);
            x->f+=x->ch[i]->f;
        }
    }
    void dfs2(nod *x){
        if(x->ed) x->g=1;
        for(int i=0;i<K;i++){
            if(!x->ch[i]) continue;
            if(!x->ch[i]->f) dfs2(x->ch[i]);
            x->f+=x->ch[i]->f;
            x->g+=x->ch[i]->g;
        }
        x->f+=x->g;
    }
    void match(){
        if(((!t)&&root->f-1<k)||(t&&root->f-root->g<k)){
            printf("-1\n");
            return ;
        }
        int rk=0;
        nod *x=root;
        while(rk<k){
            for(int i=0;i<K;i++){
                if(!x->ch[i]) continue;
                if(rk+x->ch[i]->f>=k){
                    printf("%c",i+'a');
                    x=x->ch[i];
                    if(!t) rk++;
                    else rk+=x->g;
                    break;
                }
                rk+=x->ch[i]->f;
            }
        }
    }
    void work(){
        for(nod *i=last;i->pr;i=i->pr) i->ed=1;
        if(!t) dfs1(root);
        else dfs2(root);
        match();
    }
}AM;

int main(){
    scanf("%s%d%d",s,&t,&k);
    len=strlen(s);
    for(int i=0;i<len;i++) AM.extend(s[i]-'a');
    AM.work();

    return 0;
}