BZOJ3926 [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个。

输入格式

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

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

输出格式

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

输入样例

7 3

0 2 1 2 1 0 0

1 2

3 4

3 5

4 6

5 7

2 5

输出样例

30

提示

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

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

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

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

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

题解

问题要求求出一棵树上的不同的字符串的数量，联想到后缀自动机
由于叶子结点很少，我们考虑以每一个叶子结点为根建一个trie树，那么所有的trie树一定会包含所有可能出现的字符串【作为trie树中的某一段出现，并不一定从根节点开始】
也就是说，最后大的trie树表示的所有字符串的所有子串数量，就是所求
那么用后缀自动机就可以了，与普通的SAM不同的是，这里有多个字符串一起建机，那么就是广义后缀自动机
实际上就相当于先对一个串建机，其它串再从根节点开始建机，如果已经有过对应的儿子且step[q] = step[p] + 1，那么q就已经等价于插入的节点了，就可以直接往q走，否则新建节点

求所有不同子串数量：\(\sum max[p] - max[pre[p]]\)

#include<iostream>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
#define LL long long int
#define Redge(u) for (int k = h[u],to; k; k = ed[k].nxt)
#define REP(i,n) for (int i = 1; i <= (n); i++)
#define BUG(s,n) for (int i = 1; i <= (n); i++) cout<<s[i]<<' '; puts("");
using namespace std;
const int maxn = 100005,maxm = 200005,maxv = 4000005,INF = 1000000000;
inline int read(){
	int out = 0,flag = 1; char c = getchar();
	while (c < 48 || c > 57){if (c == '-') flag = -1; c = getchar();}
	while (c >= 48 && c <= 57){out = (out << 3) + (out << 1) + c - 48; c = getchar();}
	return out * flag;
}
int h[maxn],ne = 2,de[maxn];
struct EDGE{int to,nxt;}ed[maxm];
inline void build(int u,int v){
	ed[ne] = (EDGE){v,h[u]}; h[u] = ne++;
	ed[ne] = (EDGE){u,h[v]}; h[v] = ne++;
	de[u]++; de[v]++;
}
int s[maxn],n,c;
int ch[maxv][10],pre[maxv],step[maxv],cnt;
int ins(int x,int p){
	if (step[ch[p][x]] == step[p] + 1) return ch[p][x];
	int np = ++cnt; step[np] = step[p] + 1;
	while (p && !ch[p][x]) ch[p][x] = np,p = pre[p];
	if (!p) pre[np] = 1;
	else {
		int q = ch[p][x];
		if (step[q] == step[p] + 1) pre[np] = q;
		else {
			int nq = ++cnt; step[nq] = step[p] + 1;
			for (int i = 0; i < c; i++) ch[nq][i] = ch[q][i];
			pre[nq] = pre[q]; pre[np] = pre[q] = nq;
			while (ch[p][x] == q) ch[p][x] = nq,p = pre[p];
		}
	}
	return np;
}
void dfs(int u,int last,int fa){
	int t = ins(s[u],last);
	Redge(u) if ((to = ed[k].to) != fa) dfs(to,t,u);
}
int main(){
	n = read(); c = read(); LL ans = 0;
	for (int i = 1; i <= n; i++) s[i] = read();
	for (int i = 1; i < n; i++) build(read(),read());
	cnt = 1;
	for (int i = 1; i <= n; i++)
		if (de[i] == 1) dfs(i,1,0);
	for (int i = cnt; i; i--) ans += step[i] - step[pre[i]];
	printf("%lld\n",ans);
	return 0;
}