给定一棵叶子节点不超过 \(20\) 个的无根树,每个节点上都有一个 \(0\sim 9\) 的数字,求树上本质不同路径条数。两条路径相同是指其路径上所以节点上的数字顺次连结组成的字符串相同。
\(1\le n\le 10^5\)。
如果此题是从根出发的路径,那相当于就是给定了一棵 \(\text{Trie}\) 树,非常好做。
考虑转化成从根开始的路径,对于一条路径 \(u\rightarrow v\),以 \(u\) 为根的子树内的所有节点都可以作为根将这条路径计算到。
给定的树有良好性质:叶子节点不超过 \(20\) 个。由于任意一个树上的节点的子树内必定有叶子节点,于是把每个叶子节点作为根遍历整棵树时形成的所有串加入广义后缀自动机即可。
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=2e6+5;
int n,O,col[N];char s[N];vector<int>G[N];
struct Node{
int son[10],len,fa;
Node(){memset(son,0,sizeof son);len=fa=0;}
};
struct Trie{
int cnt=1,c[N],fa[N];Node tr[N];
inline int insert(int ch,int rt){
if(!tr[rt].son[ch])tr[rt].son[ch]=++cnt,fa[cnt]=rt,c[cnt]=ch;
return tr[rt].son[ch];
}
}T1;
struct Suffix_Automaton{
int tot=1,pos[N<<1];
Node node[N<<1];
queue<int>q;
inline int add(int c,int F){
int p=F,nw=++tot;node[nw].len=node[p].len+1;
for(;p&&!node[p].son[c];p=node[p].fa)node[p].son[c]=nw;
if(!p)node[nw].fa=1;
else{
int q=node[p].son[c];
if(node[q].len==node[p].len+1)node[nw].fa=q;
else{
int xq=++tot;
node[xq]=node[q],node[xq].len=node[p].len+1;node[q].fa=node[nw].fa=xq;
for(;p&&node[p].son[c]==q;p=node[p].fa)node[p].son[c]=xq;
}
}
return nw;
}
inline void build(){
for(int i=0;i<10;++i)if(T1.tr[1].son[i])q.push(T1.tr[1].son[i]);
pos[1]=1;
while(!q.empty()){
int x=q.front();q.pop();
pos[x]=add(T1.c[x],pos[T1.fa[x]]);
for(int i=0;i<10;++i)if(T1.tr[x].son[i])q.push(T1.tr[x].son[i]);
}
}
inline ll SUM_diff(){
ll ans=0;
for(int i=2;i<=tot;++i)ans+=node[i].len-node[node[i].fa].len;
return ans;
}
}SAM;
inline void dfs(int x,int fa,int fid){
int nw=T1.insert(col[x],fid);
for(auto y:G[x])if(y^fa)dfs(y,x,nw);
}
int main(){
scanf("%d%d",&n,&O);
for(int i=1;i<=n;++i)scanf("%d",col+i);
for(int i=1,x,y;i<n;++i)scanf("%d%d",&x,&y),G[x].emplace_back(y),G[y].emplace_back(x);
for(int i=1;i<=n;++i)if((int)G[i].size()==1)dfs(i,0,1);
SAM.build();
printf("%lld\n",SAM.SUM_diff());
return 0;
}