【bzoj 3926】诸神眷顾的幻想乡（后缀自动机）

传送门biu~
题目中说叶子节点不超过20个，可以以每个叶子为根建广义后缀自动机，然后像【bzoj 4516】生成魔咒那样统计答案即可。

#include<bits/stdc++.h>
#define N 100010
using namespace std;
struct Node{
    Node *fa,*ch[10];
    int val;
    Node(int _=0){fa=NULL;memset(ch,0x0,sizeof ch);val=_;}
}*root=new Node,*point[N];
int n,c,a[N],head[N],nex[N<<1],to[N<<1],degree[N],tp;long long ans;
inline void add(int x,int y){
    nex[++tp]=head[x];
    head[x]=tp;
    to[tp]=y;
    ++degree[x],++degree[y];
}
Node* Insert(Node *p,int x){
    if(p->ch[x] && p->ch[x]->val==p->val+1)     return p->ch[x];
    Node *np=new Node(p->val+1);
    while(p && !p->ch[x])   p->ch[x]=np,p=p->fa;
    if(!p){
        np->fa=root;
        ans+=np->val;
        return np;
    }
    Node *q=p->ch[x];
    if(q->val==p->val+1){
        np->fa=q;
        ans+=np->val-q->val;
        return np;
    }
    Node *nq=new Node(p->val+1);
    nq->fa=q->fa;
    memcpy(nq->ch,q->ch,sizeof(nq->ch));
    q->fa=np->fa=nq;
    while(p && p->ch[x]==q)  p->ch[x]=nq,p=p->fa;
    ans+=np->val-np->fa->val;
    return np;
}
void dfs(int x,int father){
    point[x]=Insert(point[father],a[x]);
    for(int i=head[x];i;i=nex[i])   if(to[i]!=father) dfs(to[i],x);
}
int main(){
    scanf("%d%d",&n,&c);
    for(int i=1;i<=n;++i)   scanf("%d",&a[i]);
    for(int i=1;i<n;++i){
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    point[0]=root;
    for(int i=1;i<=n;++i)
        if(degree[i]==2)    dfs(i,0);
    printf("%lld",ans);
    return 0;
}