BZOJ 3926: [Zjoi20150]诸神眷顾的幻想乡（后缀自动机）

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
struct node{
    node* c[12];int id;
    node(){memset(c,0,sizeof(c));}
};
#define maxn 101000
vector<int> e[maxn];
#define pb push_back
int a[maxn],pre[maxn];
inline void bfs(int x,node *&y){
    if (!y) y=new node;
    static int q[maxn];
    static node * p[maxn];
    q[1]=x;
    p[1]=y;
    pre[x]=0;
    for (int l=1,r=1,u=q[1];l<=r;u=q[++l]) {
        for (vector<int>::iterator i=e[u].begin();i!=e[u].end();++i) {
            if (*i==pre[u]) continue;
            pre[*i]=u;q[++r]=*i;
            if (!p[l]->c[a[*i]]) p[l]->c[a[*i]]=new node;
            p[r]=p[l]->c[a[*i]];
        }
    }
}
struct snode{int ch[13],l,fa;}s[maxn*50];
typedef long long ll;
int cnt;ll ans;
inline void add(int x,node* u) {
    int p=++cnt,t=u->id;
    u->c[x]->id=p;
    s[p].l=s[t].l+1;
    for (;t!=-1&&!s[t].ch[x];t=s[t].fa) s[t].ch[x]=p;
    if (t==-1) s[p].fa=0;
    else if (s[t].l+1==s[s[t].ch[x]].l) s[p].fa=s[t].ch[x];
    else {
        int r=++cnt,q=s[t].ch[x];
        s[r]=s[q];s[r].l=s[t].l+1;
        s[p].fa=s[q].fa=r;
        for (;t!=-1&&s[t].ch[x]==q;t=s[t].fa) s[t].ch[x]=r;
    }
    ans+=s[p].l-s[s[p].fa].l;
}
int c;
node *root;
inline void build(){
    static node *q[maxn*50];
    q[1]=root;
    root->id=0;
    s[0].fa=-1;
    for (int l=1,r=1;l<=r;l++) {
        node *u=q[l];
        for (int i=0;i<=c;i++) 
            if (u->c[i]) {
                add(i,u);
                q[++r]=u->c[i];
            }
    }
}
int main(){
    int n;
    scanf("%d%d",&n,&c);
    for (int i=1;i<=n;i++) scanf("%d",a+i);
    for (int i=1;i<=n;i++) a[i]++;
    root=new node;
    for (int i=1;i<n;i++) {
        int x,y;
        scanf("%d%d",&x,&y);
        e[x].pb(y);e[y].pb(x);
    }
    for (int i=1;i<=n;i++) if (e[i].size()==1) bfs(i,root->c[a[i]]);
    build();
    printf("%lld\n",ans);
    return 0;
}