1 #include <iostream> 2 #include <cstdio> 3 #include <cstring> 4 #include <cmath> 5 #include <algorithm> 6 #define maxn 100005 7 #define maxl 200005 8 #define maxm 4000005 9 using namespace std; 10 11 typedef long long ll; 12 int n,c,cnt,tot,root,last,fa[maxm],du[maxn],dist[maxm],son[maxl],prep[maxl],now[maxn],v[maxn]; 13 ll ans; 14 void add(int x,int y){ 15 cnt++,prep[cnt]=now[x],now[x]=cnt,son[cnt]=y; 16 } 17 struct Tsegment{ 18 int son[maxm][10]; 19 void prepare(){tot=last=root=1;memset(dist,0,sizeof(dist));} 20 int newnode(int x){dist[++tot]=x; return tot;} 21 int add(int x,int p){ 22 int q=son[p][x]; 23 if (q==0){ 24 int np=newnode(dist[p]+1); last=np; 25 for (;p&&!son[p][x];p=fa[p]) son[p][x]=np; 26 if (p==0) fa[np]=root; 27 else{ 28 q=son[p][x]; 29 if (dist[p]+1==dist[q]) fa[np]=q; 30 else{ 31 int nq=newnode(dist[p]+1); 32 memcpy(son[nq],son[q],sizeof(son[q])); 33 fa[nq]=fa[q],fa[q]=fa[np]=nq; 34 for (;p&&son[p][x]==q;p=fa[p]) son[p][x]=nq; 35 } 36 } 37 }else{ 38 if (dist[p]+1==dist[q]) last=q; 39 else{ 40 int nq=newnode(dist[p]+1); last=nq; 41 memcpy(son[nq],son[q],sizeof(son[q])); 42 fa[nq]=fa[q],fa[q]=nq; 43 for (;p&&son[p][x]==q;p=fa[p]) son[p][x]=nq; 44 } 45 } 46 return last; 47 } 48 }SAM; 49 void dfs(int x,int pa,int goal){ 50 int p=SAM.add(v[x],goal); 51 for (int i=now[x],so=son[i];i;i=prep[i],so=son[i]){ 52 if (so==pa) continue; 53 dfs(so,x,p); 54 } 55 } 56 int main(){ 57 scanf("%d%d",&n,&c); 58 memset(du,0,sizeof(du)); 59 cnt=tot=ans=0,memset(now,0,sizeof(now)); 60 for (int i=1;i<=n;i++) scanf("%d",&v[i]); 61 for (int u,v,i=1;i<n;i++) scanf("%d%d",&u,&v),add(u,v),add(v,u),du[u]++,du[v]++; 62 SAM.prepare(); 63 for (int i=1;i<=n;i++){ 64 if (du[i]>1) continue; 65 dfs(i,0,root); 66 } 67 ans=0; 68 dist[root]=0; 69 for (int i=1;i<=tot;i++) ans+=(dist[i]-dist[fa[i]]); 70 printf("%lld ",ans); 71 return 0; 72 }
做法：广义后缀自动机模板题。
dfs+后缀自动机。