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

%%%陈老师%%%

简单来讲就是Trie上的后缀自动机，因为叶子节点总数不超过20，说以我们从每一个叶子节点出发做DFS，一边走一边把扫到的点甩进后缀自动机，这样后缀自动机里面的点不会超过4000000个，统计出本质不同的即可。

个人感觉Trie上写后缀自动机上的写法十分像树上的可持久化线段树的写法（SMG。。。先插入父亲再插入儿子

/**************************************************************
    Problem: 3926
    User: RicardoWang
    Language: C++
    Result: Accepted
    Time:2368 ms
    Memory:225364 kb
****************************************************************/
 
#include<cstdlib> 
#include<cstdio> 
#include<iostream> 
#include<cstring> 
#include<cmath> 
#include<algorithm> 
#include<queue> 
#include<vector> 
using namespace std; 
#define WJMZBMR_OTZ true 
#define maxa 4000005 
#define maxn 100005
void _read(int &x) 
{ 
    x=0; char ch=getchar(); bool flag=false; 
    while(ch<'0'||ch>'9'){if(ch=='-')flag=false;ch=getchar();} 
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0'; ch=getchar();}if(flag)x=-x; 
    return ; 
} 
 
struct edge
{
    int to,next;
}e[maxn*2];
int edge_ct,head[maxn],col[maxn],n,c,du[maxn];
void add(int x,int y)
{
    e[++edge_ct]=(edge){y,head[x]}; head[x]=edge_ct;
    e[++edge_ct]=(edge){x,head[y]}; head[y]=edge_ct; return ;
}
long long ans;
struct Sam 
{ 
    int cnt,chi[maxa][12],len[maxa],fa[maxa]; 
    Sam() 
    { 
        cnt=1;fa[1]=0;len[1]=0; 
    } 
    int extend(int last,int x) 
    { 
        int np=++cnt,p=last;len[np]=len[p]+1;
        while(!chi[p][x]&&p)chi[p][x]=np,p=fa[p]; 
        if(!p)fa[np]=1; 
        else
        { 
            int q=chi[p][x]; 
            if(len[q]==len[p]+1)fa[np]=q; 
            else
            { 
                int nq=++cnt; len[nq]=len[p]+1; 
                memcpy(chi[nq],chi[q],sizeof(chi[nq])); 
                fa[nq]=fa[q]; 
                fa[q]=fa[np]=nq; 
                while(chi[p][x]==q)chi[p][x]=nq,p=fa[p]; 
            } 
        } 
        return np; 
    } 
    void DFS(int now,int fa,int last)
    {
        int t=extend(last,col[now]);
        for(int id=head[now];id;id=e[id].next)
        {
            if(e[id].to==fa)continue;
            DFS(e[id].to,now,t);
        }
        return ;
    }
    void solve()
    {
        for(int i=1;i<=cnt;i++)
        {
            ans+=len[i]-len[fa[i]];
        }
        return ;
    }
}SAM; 
void Init()
{
    _read(n); _read(c);
    int x,y;
    for(int i=1;i<=n;i++)_read(col[i]);
    memset(du,0,sizeof(du));
    for(int i=1;i<n;i++)
    {
        _read(x); _read(y); du[x]++; du[y]++;
        add(x,y);
    }
}
int main() 
{ 
    //freopen("in.txt","r",stdin);
    Init();
    for(int i=1;i<=n;i++)if(du[i]==1)SAM.DFS(i,0,1);
    SAM.solve();
    cout<<ans<<endl;
    return 0; 
}