bzoj3926【ZJOI2015】诸神眷顾的幻想乡

3926: [Zjoi2015]诸神眷顾的幻想乡

Time Limit: 10 Sec   Memory Limit: 512 MB
Submit: 790   Solved: 485
[ Submit][ Status][ Discuss]

Description

 幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的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个。 

Input

 第一行两个正整数n,c。表示空地数量和颜色数量。 

第二行有n个0到c-1之间，由空格隔开的整数，依次表示第i块空地上的粉丝的衣服颜色。（这里我们按照节点标号从小到大的顺序依次给出每块空地上粉丝的衣服颜色）。 

接下来n-1行，每行两个正整数u,v，表示有一条连接空地u和空地v的边。 

Output

 一行，输出一个整数，表示答案。 

Sample Input

7 3
0 2 1 2 1 0 0
1 2
3 4
3 5
4 6
5 7
2 5

Sample Output

30

HINT

对于所有数据，1<=n<=100000, 1<=c<=10。 

对于15%的数据，n<=2000。 

另有5%的数据，所有空地都至多与两个空地相邻。 

另有5%的数据，除一块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。 

另有5%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻

广义后缀自动机，思路很好

题目有一个关键条件：叶子的数量不超过20个。

以每一个叶子节点为根建一个Trie树，那么整棵树中的任意一条路径，都可以表示成某一个Trie树上一个点到它后代的一条路径。

那么，可以将这些Trie树合并成一个大Trie树，然后求这个大Trie树的字串数量。

这可以用广义的后缀自动机实现。

要注意在加入一个点的时候，如果a[p][x]&&mx[a[p][x]]==mx[p]+1，则不用新建节点，直接返回a[p][x]。

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cmath>
#include<algorithm>
#define F(i,j,n) for(int i=j;i<=n;i++)
#define D(i,j,n) for(int i=j;i>=n;i--)
#define ll long long
#define maxn 100005
#define maxm 2000005
using namespace std;
int n,m,cnt,v[maxn],d[maxn],head[maxn];
ll ans;
struct edge_type{int next,to;}e[maxn*2];
struct sam
{
	int cnt,fa[maxm],mx[maxm],a[maxm][10];
	sam(){cnt=1;}	
	void solve(){F(i,1,cnt) ans+=mx[i]-mx[fa[i]];}
	int add(int p,int x)
	{
		if (a[p][x]&&mx[a[p][x]]==mx[p]+1) return a[p][x];//广义后缀自动机，注意要判断mx 
		int np=++cnt;
		mx[np]=mx[p]+1;
		while (p&&!a[p][x]) a[p][x]=np,p=fa[p];
		if (!p) fa[np]=1;
		else
		{
			int q=a[p][x];
			if (mx[p]+1==mx[q]) fa[np]=q;
			else
			{
				int nq=++cnt;
				mx[nq]=mx[p]+1;
				memcpy(a[nq],a[q],sizeof(a[q]));
				fa[nq]=fa[q];fa[np]=fa[q]=nq;
				while (a[p][x]==q) a[p][x]=nq,p=fa[p];
			}
		}
		return np;
	}
}sam;
inline int read()
{
	int x=0,f=1;char ch=getchar();
	while (ch<'0'||ch>'9'){if (ch=='-') f=-1;ch=getchar();}
	while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}
inline void add_edge(int x,int y)
{
	e[++cnt]=(edge_type){head[x],y};head[x]=cnt;
	e[++cnt]=(edge_type){head[y],x};head[y]=cnt;
}
void dfs(int x,int fa,int p)
{
	int t=sam.add(p,v[x]);
	for(int i=head[x];i;i=e[i].next) if (e[i].to!=fa) dfs(e[i].to,x,t);
}
int main()
{
	n=read();m=read();
	F(i,1,n) v[i]=read();
	F(i,1,n-1)
	{
		int x=read(),y=read();
		add_edge(x,y);
		d[x]++;d[y]++;
	}
	F(i,1,n) if (d[i]==1) dfs(i,0,1);
	sam.solve();
	printf("%lld\n",ans);
	return 0;
}