【BZOJ 3926】 [Zjoi2015]诸神眷顾的幻想乡 （广义SAM）

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<queue>
#include<cmath>
#include<set>
using namespace std;
#define Maxn 100010
#define LL long long

struct node
{
int step,pre,son[15];
}t[40*Maxn];int tot=0;
int last;

struct hp
{
int x,y,next;
}a[2*Maxn];int al=0;

int col[Maxn],fa[Maxn],first[Maxn],rd[Maxn];

void ins(int x,int y)
{
a[++al].x=x;a[al].y=y;
a[al].next=first[x];first[x]=al;
}

void upd(int x)
{
memset(t[x].son,0,sizeof(t[x].son));
t[x].pre=0;
}

int extend(int k)
{
int np=++tot;upd(tot);
t[np].step=t[last].step+1;
int now=last;
while(now&&!t[now].son[k])
{
t[now].son[k]=np;
now=t[now].pre;
}
if(!now) t[np].pre=1;
else
{
int p=now,q=t[now].son[k];
if(t[p].step+1==t[q].step) t[np].pre=q;
else
{
int nq=++tot;
upd(tot);
memcpy(t[nq].son,t[q].son,sizeof(t[nq].son));
t[nq].pre=t[q].pre;
t[q].pre=t[np].pre=nq;
t[nq].step=t[p].step+1;
while(now&&t[now].son[k]==q)
{
t[now].son[k]=nq;
now=t[now].pre;
}
}
}
last=np;
return np;
}

void dfs(int x,int f)
{
int now=extend(col[x]);
for(int i=first[x];i;i=a[i].next) if(a[i].y!=f)
{
last=now;
dfs(a[i].y,x);
}
}

int n,c;
void init()
{
scanf("%d%d",&n,&c);
for(int i=1;i<=n;i++) {scanf("%d",&col[i]);col[i]++;}
memset(first,0,sizeof(first));
memset(rd,0,sizeof(rd));
for(int i=1;i<n;i++)
{
int x,y;
scanf("%d%d",&x,&y);
rd[x]++;rd[y]++;
ins(x,y);ins(y,x);
}
t[++tot].step=0;
upd(tot);last=1;
for(int i=1;i<=n;i++) if(rd[i]==1)
{
dfs(i,0);
last=1;
}
}

int main()
{
init();
LL ans=0;
for(int i=1;i<=tot;i++) ans+=t[i].step-t[t[i].pre].step;
printf("%lld\n",ans);
return 0;
}