[BZOJ3926][Zjoi2015]诸神眷顾的幻想乡

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

题解

叶子不超过 20 个。。。看成度数不超过 20 了。。。。。

于是我们可以以每个叶子为根搞出一颗 trie 树，于是得到不超过 20 颗 trie 树，然后合并这些 trie 树，在上面构造后缀自动机，然后就是裸题了。

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cctype>
#include <algorithm>
using namespace std;

int read() {
	int x = 0, f = 1; char c = getchar();
	while(!isdigit(c)){ if(c == '-') f = -1; c = getchar(); }
	while(isdigit(c)){ x = x * 10 + c - '0'; c = getchar(); }
	return x * f;
}

#define maxn 100010
#define maxm 200010
#define maxnode 2000010
#define maxNode 4000010
#define maxa 10
#define LL long long

int n, m, head[maxn], nxt[maxm], To[maxm], col[maxn], ind[maxn];

void AddEdge(int a, int b) {
	To[++m] = b; nxt[m] = head[a]; head[a] = m;
	swap(a, b);
	To[++m] = b; nxt[m] = head[a]; head[a] = m;
	ind[a]++; ind[b]++;
	return ;
}

int tot, rt, ch[maxnode][maxa], fa[maxnode], fac[maxnode];
void dfs(int u, int p, int pa) {
	if(!ch[p][col[u]]) ch[p][col[u]] = ++tot, fa[tot] = p, fac[tot] = col[u];
	for(int e = head[u]; e; e = nxt[e]) if(To[e] != pa)
		dfs(To[e], ch[p][col[u]], u);
	return ;
}

int Rt, ToT, to[maxNode][maxa], par[maxNode], Max[maxNode], last[maxnode];
void extend(int x, int p, int id) {
	int np = ++ToT; Max[np] = Max[p] + 1; last[id] = np;
	while(p && !to[p][x]) to[p][x] = np, p = par[p];
	if(!p){ par[np] = rt; return ; }
	int q = to[p][x];
	if(Max[q] == Max[p] + 1){ par[np] = q; return ; }
	int nq = ++ToT; Max[nq] = Max[p] + 1;
	memcpy(to[nq], to[q], sizeof(to[q]));
	par[nq] = par[q];
	par[q] = par[np] = nq;
	while(p && to[p][x] == q) to[p][x] = nq, p = par[p];
	return ;
}

int main() {
	n = read(); read();
	for(int i = 1; i <= n; i++) col[i] = read();
	for(int i = 1; i < n; i++) {
		int a = read(), b = read();
		AddEdge(a, b);
	}
	
	tot = rt = 1;
	for(int i = 1; i <= n; i++) if(ind[i] == 1)
		dfs(i, rt, 0);
	Rt = ToT = 1; last[1] = 1;
	for(int i = 2; i <= tot; i++) extend(fac[i], last[fa[i]], i);
//	for(int i = 1; i <= ToT; i++) printf("%d(%d)%c", par[i], i, i < ToT ? ' ' : '
');
//	for(int i = 1; i <= ToT; i++) printf("%d(%d)%c", rgt[i], i, i < ToT ? ' ' : '
');
//	for(int i = 1; i <= ToT; i++) printf("%d(%d)%c", Max[i], i, i < ToT ? ' ' : '
');
	LL ans = 0;
	for(int i = 1; i <= ToT; i++) ans += Max[i] - Max[par[i]];
	
	printf("%lld
", ans);
	
	return 0;
}