ZJOI2015 诸神眷顾的幻想乡

诸神眷顾的幻想乡

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的 2600 岁生日，无数幽香的粉丝到了幽香家门前的太阳花田上来为幽香庆祝生日。
 
粉丝们非常热情，自发组织表演了一系列节目给幽香看。幽香当然也非常高兴啦。
 
这时幽香发现了一件非常有趣的事情，太阳花田有 \(n\) 块空地。在过去，幽香为了方便，在这 \(n\) 块空地之间修建了 \(n-1\) 条边将它们连通起来。也就是说，这 \(n\) 块空地形成了一个树的结构。

有 \(n\) 个粉丝们来到了太阳花田上。为了表达对幽香生日的祝贺，他们选择了 \(c\) 种颜色的衣服，每种颜色恰好可以用一个 \(0\) 到 \(c-1\) 之间的整数来表示。并且每个人都站在一个空地上，每个空地上也只有一个人。这样整个太阳花田就花花绿绿了。幽香看到了，感觉也非常开心。
 
粉丝们策划的一个节目是这样的，选中两个粉丝 \(A\) 和 \(B\)（\(A\) 和 \(B\) 可以相同），然后 \(A\) 所在的空地到 \(B\) 所在的空地的路径上的粉丝依次跳起来（包括端点），幽香就能看到一个长度为 \(A\) 到 \(B\) 之间路径上的所有粉丝的数目（包括 \(A\) 和 \(B\)）的颜色序列。一开始大家打算让每两个粉丝（注意：\(A,B\) 和 \(B,A\) 是不同的，他们形成的序列刚好相反，比如红绿蓝和蓝绿红）都来一次，但是有人指出这样可能会出现一些一模一样的颜色序列，会导致审美疲劳。

于是他们想要问题，在这个树上，一共有多少可能的不同的颜色序列（子串）幽香可以看到呢？

太阳花田的结构比较特殊，只与一个空地相邻的空地数量不超过 \(20\) 个。

对于所有数据，\(1 \leq n \leq 100000, 1 \leq c \leq 10\)。

分析

广义SAM的板子题。

首先叶子节点不超过20，那么可以直接对每个叶子节点为根的子树插入到广义SAM中。

因为所有合法的答案一定是某个叶子节点为根的树上的一条链，因此这样可以统计出所有合法的答案

然后就是经典的本质不同子串问题了，\(ans = \sum (len[i] - len[fa[i]])\)

时间复杂度\(O(20*n)\)

co int N=4e6;
int tot=1;
int ch[N][10],fa[N],len[N];
int extend(int p,int c){
	int last;
	if(ch[p][c]){
		int q=ch[p][c];
		if(len[q]==len[p]+1) last=q;
		else{
			int clone=last=++tot;
			memcpy(ch[clone],ch[q],sizeof ch[q]);
			fa[clone]=fa[q],len[clone]=len[p]+1;
			fa[q]=clone;
			for(;ch[p][c]==q;p=fa[p]) ch[p][c]=clone;
		}
	}
	else{
		int cur=last=++tot;
		len[cur]=len[p]+1;
		for(;p&&!ch[p][c];p=fa[p]) ch[p][c]=cur;
		if(!p) fa[cur]=1;
		else{
			int q=ch[p][c];
			if(len[q]==len[p]+1) fa[cur]=q;
			else{
				int clone=++tot;
				memcpy(ch[clone],ch[q],sizeof ch[q]);
				fa[clone]=fa[q],len[clone]=len[p]+1;
				fa[cur]=fa[q]=clone;
				for(;ch[p][c]==q;p=fa[p]) ch[p][c]=clone;
			}
		}
	}
	return last;
}
int n,col[N],deg[N];
vector<int> edg[N];
void dfs(int u,int fa,int last){
	int new_last=extend(last,col[u]);
	for(int i=0,v;i<edg[u].size();++i){
		if((v=edg[u][i])==fa) continue;
		dfs(v,u,new_last);
	}
}
int main(){
	read(n),read<int>();
	for(int i=1;i<=n;++i) read(col[i]);
	for(int i=1,u,v;i<n;++i){
		read(u),read(v);
		edg[u].push_back(v),edg[v].push_back(u);
		++deg[u],++deg[v];
	}
	for(int i=1;i<=n;++i)
		if(deg[i]==1) dfs(i,0,1);
	ll ans=0;
	for(int i=1;i<=tot;++i) ans+=len[i]-len[fa[i]];
	printf("%lld\n",ans);
	return 0;
}