BZOJ3926 ZJOI2015 诸神眷顾的幻想乡 后缀自动机+DFS

#include <queue>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <climits>
#include <iostream>
#include <algorithm>
using namespace std;
#define ll long long

const int MAXC=10+2;
const int MAXN=100000+2;
struct HASH{
    int u;
    HASH *next;
    HASH(){}
    HASH(int _u,HASH *_next):u(_u),next(_next){}
}*table[MAXN],mem[2*MAXN];
struct SAM{
    int v;
    SAM *child[MAXC],*f;
    bool flag;
    SAM(){}
    SAM(int _v):v(_v),f(0),flag(0){ memset(child,0,sizeof(child));}
}*root,*last=root=new SAM(0),*node[MAXN];
int N,C,c[MAXN],d[MAXN],cnt;
queue<SAM *> q;

void Insert(int u,int v){ table[u]=&(mem[cnt++]=HASH(v,table[u]));}

SAM *Extend(SAM *&x,int c){
    if(x->child[c] && x->child[c]->v==x->v+1) return x->child[c];

    SAM *p=x,*np=new SAM(p->v+1);
    while(p && !p->child[c]) p->child[c]=np,p=p->f;

    if(!p) np->f=root;
    else{
        SAM *q=p->child[c];
        if(q->v==p->v+1) np->f=q;
        else{
            SAM *nq=new SAM(p->v+1);
            memcpy(nq->child,q->child,sizeof(q->child));
            nq->f=q->f,q->f=np->f=nq;
            while(p && p->child[c]==q) p->child[c]=nq,p=p->f;
        }
    }
    return np;
}

void DFS(int x,int f){
    node[x]=Extend(node[f],c[x]);
    for(HASH *p=table[x];p;p=p->next)
        if(p->u!=f) DFS(p->u,x);
}

ll BFS(SAM *x){
    ll ret=0;
    q.push(x);

    SAM *t;
    while(!q.empty()){
        t=q.front(),q.pop();
        if(t!=x) ret+=t->v-t->f->v;
        for(int i=0;i<C;i++)
            if(t->child[i] && !t->child[i]->flag)
                t->child[i]->flag=1,q.push(t->child[i]);
    }
    return ret;
}

int main(){
    scanf("%d %d",&N,&C);
    for(int i=1;i<=N;i++) scanf("%d",c+i);
    for(int i=1,a,b;i<N;i++){
        scanf("%d %d",&a,&b);
        d[a]++,d[b]++;
        Insert(a,b),Insert(b,a);
    }

    node[0]=root;
    for(int i=1;i<=N;i++)
        if(d[i]==1) DFS(i,0);

    printf("%lld\n",BFS(root));

    return 0;
}