【BZOJ3926】诸神眷顾的幻想乡 【广义后缀自动机】

题意

给定一棵树，每个结点有一个颜色，问树上有多少种子串（定义子串为两点上路径颜色的序列）。保证叶子结点<=20

分析

我们可以发现一个结论，任意一个子串一定是以某个叶子结点为根的trie的后缀。我们有注意到，叶子节点最多只有20，那么我们可以将每个叶子结点拿出来，以它为根按照trie树的方式插到广义后缀自动机中。要统计一共有多少子串的话，这样这样建好自动机以后枚举每个状态，然后res+=st[u].len-st[st[u].link].len;

 #include <cstdio>
 #include <cstring>
 #include <algorithm>
 #include <iostream>

 using namespace std;
 const int maxn=+;
 typedef long long LL;
 int val[maxn];
 int head[maxn],Next[*maxn],to[*maxn],deg[maxn];
 int N,sz,C;
 int cnt,cur;
 struct state{
     int len,link;
     int next[];
 }st[*maxn];

 void init(){
     sz=;
     memset(head,-,sizeof(head));
     cnt=;
     cur=;
     st[].link=-;
     st[].len=;
 }

 void add_edge(int a,int b){
     ++sz;
     to[sz]=b;
     Next[sz]=head[a];
     head[a]=sz;
 }

 int build_sam(int c,int last){
     cur=cnt++;
     st[cur].len=st[last].len+;
     int p;
     for(p=last;p!=-&&st[p].next[c]==;p=st[p].link)
         st[p].next[c]=cur;
     if(p==-)
         st[cur].link=;
     else{
         int q=st[p].next[c];
         if(st[q].len==st[p].len+)
             st[cur].link=q;
         else{
             int clone=cnt++;
             st[clone].len=st[p].len+;
             st[clone].link=st[q].link;
             for(int i=;i<C;i++)
                 st[clone].next[i]=st[q].next[i];
             for(;p!=-&&st[p].next[c]==q;p=st[p].link)
                 st[p].next[c]=clone;
             st[cur].link=st[q].link=clone;
         }
     }
     return cur;
 }

 void dfs(int u,int fa,int las){
     int Las=build_sam(val[u],las);
     for(int i=head[u];i!=-;i=Next[i]){
         int v=to[i];
         if(v==fa)continue;
         dfs(v,u,Las);
     }
     return;
 }

 int main(){
     scanf("%d%d",&N,&C);
     for(int i=;i<=N;i++)
         scanf("%d",&val[i]);
     int a,b;
     init();
     for(int i=;i<N;i++){
         scanf("%d%d",&a,&b);
         add_edge(a,b);
         add_edge(b,a);
         deg[a]++;deg[b]++;
     }
     for(int i=;i<=N;i++){
         if(deg[i]==){
             dfs(i,-,);
         }
     }
     LL ans=;
     for(int i=;i<cnt;i++){
         int p=st[i].link;
         if(p!=-){
             ans+=st[i].len-st[p].len;
         }
     }
     printf("%lld\n",ans);
 return ;
 }