听说这是广义(SAM)的板子
看来对于广义(SAM)我也就只会板子了
叶子数很少,所以可以枚举每一个叶子节点作为根建一遍(Trie)树
只需要对(Trie)树建出(SAM)就好了
跟对单串建(SAM)不同的就是(last)节点是这个点在(Trie)树上的父亲
并不是很清楚为什么需要在(son[f][c]!=0)的时候特判一下
就这样吧
#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#define maxn 100005
#define M 4000005
#define re register
#define LL long long
#define max(a,b) ((a)>(b)?(a):(b))
#define min(a,b) ((a)<(b)?(a):(b))
inline int read()
{
char c=getchar();int x=0;while(c<'0'||c>'9') c=getchar();
while(c>='0'&&c<='9') x=(x<<3)+(x<<1)+c-48,c=getchar();
return x;
}
struct E{int v,nxt;} e[maxn<<1];
int c[maxn],head[maxn],lst=1,cnt=1,d[maxn],num,n,m;
int len[M],fa[M],son[M][10];
inline void add(int x,int y) {e[++num].v=y;e[num].nxt=head[x];head[x]=num;}
inline int ins(int f,int c)
{
if(son[f][c])
{
int x=son[f][c];
if(len[f]+1==len[x]) {return x;}
int y=++cnt;
len[y]=len[f]+1,fa[y]=fa[x],fa[x]=y;
for(re int i=0;i<m;i++) son[y][i]=son[x][i];
while(f&&son[f][c]==x) son[f][c]=y,f=fa[f];
return y;
}
int p=++cnt; lst=p;
len[p]=len[f]+1;
while(f&&!son[f][c]) son[f][c]=p,f=fa[f];
if(!f) {fa[p]=1;return lst;}
int x=son[f][c];
if(len[f]+1==len[x]) {fa[p]=x;return lst;}
int y=++cnt;
len[y]=len[f]+1,fa[y]=fa[x],fa[x]=fa[p]=y;
for(re int i=0;i<m;i++) son[y][i]=son[x][i];
while(f&&son[f][c]==x) son[f][c]=y,f=fa[f];
return lst;
}
void dfs(int x,int f,int p)
{
int t=ins(p,c[x]);
for(re int i=head[x];i;i=e[i].nxt)
if(e[i].v!=f) dfs(e[i].v,x,t);
}
int main()
{
n=read(),m=read();int x,y;LL ans=0;
for(re int i=1;i<=n;i++) c[i]=read();
for(re int i=1;i<n;i++) x=read(),y=read(),add(x,y),add(y,x),d[x]++,d[y]++;
for(re int i=1;i<=n;i++) if(d[i]==1) dfs(i,0,1);
for(re int i=2;i<=cnt;i++) ans+=len[i]-len[fa[i]];
printf("%lld
",ans);
return 0;
}