Luogu3346[ZJOI2015]诸神眷顾的幻想乡
题面:洛谷
解析
观察到树的叶子节点数量很少(题面描述是真的迷,我开始以为是度数小于20),发现对于树上的每一条路径,都能在以某一个叶子节点为根的树中表示为一条从上到下的路径,那么把每一颗树看做一颗(Trie)树,用广义后缀自动机插入即可,答案就是不同状态与他后缀链接的状态(len)之差的和。
代码
// luogu-judger-enable-o2
#include<cstdio>
#include<cstring>
#define N 100005
#define LL long long
using namespace std;
int n,c,col[N];
inline int In(){
char c=getchar(); int x=0,ft=1;
for(;c<'0'||c>'9';c=getchar()) if(c=='-') ft=-1;
for(;c>='0'&&c<='9';c=getchar()) x=x*10+c-'0';
return x*ft;
}
namespace Sam{
int root,sz,ch[N*20][10],fa[N*20],len[N*20];
inline void init(){
sz=0; root=++sz;
}
inline int newnode(int u){
len[++sz]=len[u]+1; return sz;
}
inline int Extend(int f,int c){
if(ch[f][c]&&len[ch[f][c]]==len[f]+1) return ch[f][c];
int p=newnode(f);
while(f&&!ch[f][c]) ch[f][c]=p,f=fa[f];
if(!f){ fa[p]=root; return p; }
int x=ch[f][c],ff=0;
if(len[x]==len[f]+1){ fa[p]=x; return p; }
if(len[p]==len[f]+1) ff=1;
int y=newnode(f); memcpy(ch[y],ch[x],sizeof(ch[y]));
fa[y]=fa[x]; fa[x]=fa[p]=y;
while(f&&ch[f][c]==x) ch[f][c]=y,f=fa[f];
return ff?y:p;
}
inline void Calc(){
LL ans=0;
for(int i=root;i<=sz;++i) ans+=len[i]-len[fa[i]];
printf("%lld
",ans);
}
}
int h[N],deg[N],e_tot=0;
struct E{ int to,nex; }e[N<<1];
inline void add(int u,int v){
e[++e_tot]=(E){v,h[u]}; h[u]=e_tot; ++deg[u];
e[++e_tot]=(E){u,h[v]}; h[v]=e_tot; ++deg[v];
}
void dfs(int u,int las,int fa){
int v_las=Sam::Extend(las,col[u]);
for(int i=h[u],v;i;i=e[i].nex){
v=e[i].to; if(v==fa) continue;
dfs(v,v_las,u);
}
}
int main(){
n=In(); c=In(); Sam::init();
for(int i=1;i<=n;++i) col[i]=In();
for(int i=1;i<n;++i) add(In(),In());
for(int i=1;i<=n;++i) if(deg[i]==1) dfs(i,Sam::root,-1);
Sam::Calc();
return 0;
}