BZOJ 2894: 世界线 广义后缀自动机

#include<bits/stdc++.h>
#define maxn 300000 
#define ll long long 
using namespace std; 
vector<int>G[maxn]; 
char str[maxn];   
int n,Q; 
ll ans=1; 
void setIO(string s)
{
    string in=s+".in"; 
    string out=s+".out"; 
    freopen(in.c_str(),"r",stdin); 
    freopen(out.c_str(),"w",stdout); 
}       
namespace SAM
{
    int tot; 
    int ch[maxn<<1][28], f[maxn<<1], len[maxn<<1], cnt[maxn<<1], sumv[maxn<<1], C[maxn<<1], rk[maxn<<1];   
    char tr[maxn];                 
    void init()
    {
        tot=1; 
    }
    int extend(int c,int last)
    {
        int p=last; 
        if(!ch[p][c])
        {
            int np=++tot; 
            last=np, len[np]=len[p]+1; 
            while(p&&!ch[p][c])ch[p][c]=np,p=f[p]; 
            if(!p) f[np]=1; 
            else 
            {
                int q=ch[p][c]; 
                if(len[q]==len[p]+1)  f[np]=q; 
                else 
                {
                    int nq=++tot; 
                    len[nq]=len[p]+1; 
                    memcpy(ch[nq],ch[q],sizeof(ch[q])); 
                    f[nq]=f[q], f[np]=f[q]=nq; 
                    while(p&&ch[p][c]==q) ch[p][c]=nq, p=f[p];    
                } 
            }  
            ans+=len[np]-len[f[np]]; 
        }
        else 
        {
            int q=ch[p][c];
            if(len[q]==len[p]+1) last=q;               
            else     
            {
                int nq=++tot; 
                len[nq]=len[p]+1; 
                memcpy(ch[nq],ch[q],sizeof(ch[q])); 
                f[nq]=f[q], f[q]=nq; 
                while(p&&ch[p][c]==q) ch[p][c]=nq,p=f[p];         
                last=nq;        
            }
        }         
        return last;  
    }     
    void Get()
    {
        for(int i=1;i<=tot;++i) ++C[len[i]]; 
        for(int i=1;i<=tot;++i) C[i]+=C[i-1]; 
        for(int i=1;i<=tot;++i) rk[C[len[i]]--]=i; 
        for(int i=tot;i>=1;--i)
        {
            int p=rk[i];   
            sumv[p]=1;        
            for(int j=0;j<26;++j) if(ch[p][j]) sumv[p]+=sumv[ch[p][j]]; 
        }
    }
    void solve(int k)
    {
        int cur=1,cc=0,flag2=1;  
        while(flag2 && cur)
        {                
            flag2=0;                
            if(cur!=1) k-=1;     
            if(k<=0)          
            {
                for(int i=1;i<=cc;++i) printf("%c",tr[i]);           
                puts(""); 
                return; 
            }         
            for(int i=0;i<27;++i)
            {
                if(ch[cur][i]) 
                {   
                    if(sumv[ch[cur][i]] >= k) 
                    { 
                        cur=ch[cur][i];
                        tr[++cc]=i+‘a‘; 
                        flag2=1;   
                        break;
                    } 
                    else k-=sumv[ch[cur][i]];        
                }
            }          
        }
        puts("-1"); 
    }
}; 
void DFS(int u,int fa,int cur)
{        
    cur=SAM::extend(str[u]-‘a‘,cur);   // 插入该字符后的状态           
    for(int i=0;i<G[u].size();++i)
    { 
        int v=G[u][i]; 
        if(v==fa) continue;   
        DFS(v,u,cur);     
    }
}
int main()
{
    // setIO("input"); 
    scanf("%d%d",&n,&Q); 
    scanf("%s",str+1); 
    for(int i=1,a,b;i<n;++i)
    {
        scanf("%d%d",&a,&b); 
        G[a].push_back(b); 
        G[b].push_back(a); 
    }
    SAM::init(); 
    DFS(1,0,1);    
    SAM::Get();     
    printf("%lld",ans);      
    puts(""); 
    while(Q--)
    {
        int k; 
        scanf("%d",&k); --k; 
        if(k) SAM::solve(k);  
        else puts(""); 
    } 
    return 0; 
}