P3346 [ZJOI2015]诸神眷顾的幻想乡

$ \color{#0066ff}{ 题目描述 }$

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

\(\color{#0066ff}{输入格式}\)

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

\(\color{#0066ff}{输出格式}\)

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

\(\color{#0066ff}{输入样例}\)

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

\(\color{#0066ff}{输出样例}\)

30

\(\color{#0066ff}{数据范围与提示}\)

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

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

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

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

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

\(\color{#0066ff}{题解}\)

显然题目是找不同字串个数，如果给我们的树是一条链，那么直接SAM就水了

现在是一棵树，如果用后缀自动机的话，怎么把所有的极长链都扔进去呢?

不难发现， 极长链的两端一定度为1，题目还告诉你度为1的点不超过20，就可以暴力了

暴力把所有的极长链正反搜出来，加入SAM，统计就行了

#include<bits/stdc++.h>
#define LL long long
LL in() {
	char ch; LL x = 0, f = 1;
	while(!isdigit(ch = getchar()))(ch == '-') && (f = -f);
	for(x = ch ^ 48; isdigit(ch = getchar()); x = (x << 1) + (x << 3) + (ch ^ 48));
	return x * f;
}
const int maxn = 2e6 + 100;
struct node {
	node *ch[13], *fa;
	int len, siz;
	node(int len = 0, int siz = 0): len(len), siz(siz) {
		fa = NULL; memset(ch, 0, sizeof ch); 
	}
}pool[maxn], *tail, *root;
node *extend(int c, node *lst) {
	node *o = new(tail++) node(lst->len + 1, 1), *v = lst;
	for(; v && !v->ch[c]; v = v->fa) v->ch[c] = o;
	if(!v) o->fa = root;	
	else if(v->len + 1 == v->ch[c]->len) o->fa = v->ch[c];
	else {
		node *n = new(tail++) node(v->len + 1), *d = v->ch[c];
		std::copy(d->ch, d->ch + 13, n->ch);
		n->fa = d->fa, d->fa = o->fa = n;
		for(; v && v->ch[c] == d; v = v->fa) v->ch[c] = n;
	}
	return o;
}
void init() {
	tail = pool;
	root = new(tail++) node();
}
int du[maxn], n, c, val[maxn];
std::vector<int> G[maxn];
void dfs(int x, int fa, node *last) {
	node *now = extend(val[x], last);
	for(int to : G[x]) if(to != fa) dfs(to, x, now);
}
int main() {
	n = in(), c = in();
	for(int i = 1; i <= n; i++) val[i] = in();
	int x, y;
	for(int i = 1; i < n; i++) {
		du[x = in()]++, du[y = in()]++;
		G[x].push_back(y), G[y].push_back(x);
	}
	init();
	for(int i = 1; i <= n; i++) if(du[i] == 1) dfs(i, 0, root);
	LL ans = 0;
	for(node *o = pool + 1; o != tail; o++) ans += o->len - o->fa->len;
	printf("%lld\n", ans);
	return 0;
}