题意
给定一棵树,每个结点有一个颜色,问树上有多少种子串(定义子串为两点上路径颜色的序列)。保证叶子结点<=20
分析
我们可以发现一个结论,任意一个子串一定是以某个叶子结点为根的trie的后缀。我们有注意到,叶子节点最多只有20,那么我们可以将每个叶子结点拿出来,以它为根按照trie树的方式插到广义后缀自动机中。要统计一共有多少子串的话,这样这样建好自动机以后枚举每个状态,然后res+=st[u].len-st[st[u].link].len;
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream> using namespace std;
const int maxn=+;
typedef long long LL;
int val[maxn];
int head[maxn],Next[*maxn],to[*maxn],deg[maxn];
int N,sz,C;
int cnt,cur;
struct state{
int len,link;
int next[];
}st[*maxn]; void init(){
sz=;
memset(head,-,sizeof(head));
cnt=;
cur=;
st[].link=-;
st[].len=;
} void add_edge(int a,int b){
++sz;
to[sz]=b;
Next[sz]=head[a];
head[a]=sz;
} int build_sam(int c,int last){
cur=cnt++;
st[cur].len=st[last].len+;
int p;
for(p=last;p!=-&&st[p].next[c]==;p=st[p].link)
st[p].next[c]=cur;
if(p==-)
st[cur].link=;
else{
int q=st[p].next[c];
if(st[q].len==st[p].len+)
st[cur].link=q;
else{
int clone=cnt++;
st[clone].len=st[p].len+;
st[clone].link=st[q].link;
for(int i=;i<C;i++)
st[clone].next[i]=st[q].next[i];
for(;p!=-&&st[p].next[c]==q;p=st[p].link)
st[p].next[c]=clone;
st[cur].link=st[q].link=clone;
}
}
return cur;
} void dfs(int u,int fa,int las){
int Las=build_sam(val[u],las);
for(int i=head[u];i!=-;i=Next[i]){
int v=to[i];
if(v==fa)continue;
dfs(v,u,Las);
}
return;
} int main(){
scanf("%d%d",&N,&C);
for(int i=;i<=N;i++)
scanf("%d",&val[i]);
int a,b;
init();
for(int i=;i<N;i++){
scanf("%d%d",&a,&b);
add_edge(a,b);
add_edge(b,a);
deg[a]++;deg[b]++;
}
for(int i=;i<=N;i++){
if(deg[i]==){
dfs(i,-,);
}
}
LL ans=;
for(int i=;i<cnt;i++){
int p=st[i].link;
if(p!=-){
ans+=st[i].len-st[p].len;
}
}
printf("%lld\n",ans);
return ;
}