Luogu-3346 [ZJOI2015]诸神眷顾的幻想乡

\(trie\)树建广义后缀自动机：

\(dfs\)遍历\(trie\)树，将树上的一个节点插入\(sam\)时，将他的\(fa\)在\(sam\)上所在的节点作为\(last\)

#include<map>
#include<queue>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
const int maxn=2e5+100,maxm=4e6+100;
struct SAM{
    int son[maxm][11],len[maxm],fa[maxm];
    int last,tot,head[maxn],num,nex[maxn],v[maxn],c[maxn],du[maxn];
    SAM(){last=tot=0,fa[0]=-1,num=1;}
    void add(int x,int y){
        v[++num]=y;
        nex[num]=head[x];
        head[x]=num;
        v[++num]=x;
        nex[num]=head[y];
        head[y]=num;
        du[x]++,du[y]++;
    }
    int insert(int x,int last){
        int p=last,np=son[p][x];
        if(np&&len[np]==len[p]+1) return np;
        np=++tot;
        len[np]=len[p]+1;
        while(~p&&!son[p][x])
            son[p][x]=np,p=fa[p];
        if(p==-1)
            fa[np]=0;
        else{
            int q=son[p][x];
            if(len[q]==len[p]+1)
                fa[np]=q;
            else{
                int nq=++tot;
                memcpy(son[nq],son[q],sizeof(son[q]));
                fa[nq]=fa[q];
                len[nq]=len[p]+1;
                fa[q]=fa[np]=nq;
                while(~p&&son[p][x]==q)
                    son[p][x]=nq,p=fa[p];
            }
        }
        return np;
    }
    void dfs(int x,int fa,int last){
        last=insert(c[x],last);
        for(int i=head[x];i;i=nex[i])
            if(v[i]!=fa)
                dfs(v[i],x,last);
    }
    void query(){
        ll ans=0;
        for(int i=1;i<=tot;i++)
            ans+=1ll*(len[i]-len[fa[i]]);
        printf("%lld\n",ans);
    }
}sam;
int n,m,a,b;
int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++) scanf("%d",&sam.c[i]);
    for(int i=1;i<n;i++)
        scanf("%d%d",&a,&b),sam.add(a,b);
    for(int i=1;i<=n;i++)
        if(sam.du[i]==1)
            sam.dfs(i,i,0);
    sam.query();
    return 0;
}