bzoj3926: [Zjoi2015]诸神眷顾的幻想乡

#include<cstdio>
#include<iostream>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<cmath>
using namespace std;
typedef long long LL;

struct node
{
    int x,y,next;
}a[210000];int len,last[110000];
void ins(int x,int y)
{
    len++;
    a[len].x=x;a[len].y=y;
    a[len].next=last[x];last[x]=len;
}
int lf[30];
void findleaf(int x,int fr)
{
    int t=0;
    for(int k=last[x];k;k=a[k].next)
    {
        int y=a[k].y;
        if(y!=fr)
        {
            t++;
            findleaf(y,x);
        }
    }
    if(t==0|(x==1&&t==1))lf[++lf[0]]=x;
}

//---------------------------------init:findleaf-----------------------------------------------------

int col[110000];
struct Trie
{
    int w[15],col;
    void init(int COL){memset(w,0,sizeof(w));col=COL;}
}tr[2100000];int trlen;
int Trie_insert(int x,int now)
{
    if(tr[now].w[x]==0)
        tr[now].w[x]=++trlen, tr[trlen].init(x);
    return tr[now].w[x];
}

int fr[2100000],pos[2100000];
int head,tail,list[2100000];
void Trie_merge(int rt)
{
    fr[rt]=0,pos[0]=0;
    head=1,tail=2;list[1]=rt;
    while(head!=tail)
    {
        int x=list[head]; pos[x]=Trie_insert(col[x],pos[fr[x]]);
        for(int k=last[x];k;k=a[k].next)
        {
            int y=a[k].y;
            if(y!=fr[x])
            {
                fr[y]=x;
                list[tail++]=y;
            }
        }
        fr[x]=0;
        head++;
    }
}
void Trie_main()
{
    trlen=0;
    for(int i=1;i<=lf[0];i++)Trie_merge(lf[i]);
}

//---------------------------------------maketrie---------------------------------------------------

struct SAM
{
    int w[15],dep,fail;
}ch[4100000];int cnt;
int SAM_insert(int pre,int x)
{
    int now=++cnt;
    ch[now].dep=ch[pre].dep+1;
    
    while(pre!=0&&ch[pre].w[x]==0)
        ch[pre].w[x]=now, pre=ch[pre].fail;
        
    if(pre==0)ch[now].fail=1;
    else
    {
        int nxt=ch[pre].w[x];
        if(ch[nxt].dep==ch[pre].dep+1)ch[now].fail=nxt;
        else
        {
            int nnxt=++cnt;
            ch[nnxt]=ch[nxt];
            ch[nnxt].dep=ch[pre].dep+1;
            
            ch[nxt].fail=ch[now].fail=nnxt;
            while(pre!=0&&ch[pre].w[x]==nxt)
                ch[pre].w[x]=nnxt, pre=ch[pre].fail;
        }
    }
    return now;
}

int C;
void SAM_main()
{
    cnt=1;
    head=1,tail=1;
    pos[0]=1;
    for(int x=1;x<=C;x++)
        if(tr[0].w[x]!=0) fr[tr[0].w[x]]=0, list[tail++]=tr[0].w[x];
    while(head!=tail)
    {
        int now=list[head]; pos[now]=SAM_insert(pos[fr[now]],tr[now].col);
        for(int x=1;x<=C;x++)
            if(tr[now].w[x]!=0)
            {
                fr[tr[now].w[x]]=now;
                list[tail++]=tr[now].w[x];
            }
        head++;
    }
}

//---------------------------------------makeSAM----------------------------------------------------

int main()
{
    freopen("a.in","r",stdin);
    freopen("a.out","w",stdout);
    int n,x,y;
    scanf("%d%d",&n,&C);
    len=0;memset(last,0,sizeof(last));
    for(int i=1;i<=n;i++)scanf("%d",&col[i]),col[i]++;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&x,&y);
        ins(x,y),ins(y,x);
    }
    findleaf(1,0);
    
    Trie_main();
    SAM_main();
    
    LL ans=0;
    for(int i=2;i<=cnt;i++)
        ans+=ch[i].dep-ch[ch[i].fail].dep;
    printf("%lld\n",ans);
    
    return 0;
}