「BZOJ3926」[ZJOI2015]诸神眷顾的幻想乡

Description

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

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

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

太阳花田的结构比较特殊,只与一个空地相邻的空地数量不超过 2020 个。

Input

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

第二行有 nn00c1c-1之间,由空格隔开的整数,依次表示第 ii 块空地上的粉丝的衣服颜色。(这里我们按照节点标号从小到大的顺序依次给出每块空地上粉丝的衣服颜色)。

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

Output

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

Sample Input

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

Sample Output

1
30

HINT

对于所有数据,1n100000,1c101 \leq n \leq 100000, 1 \leq c \leq 10

对于 15%15 \% 的数据,n2000n \leq 2000
 
另有 5%5 \% 的数据,所有空地都至多与两个空地相邻。
 
另有 5%5 \% 的数据,除一块空地与三个空地相邻外,其他空地都分别至多与两个空地相邻。
 
另有 5%5 \% 的数据,除某两块空地与三个空地相邻外,其他空地都分别至多与两个空地相邻。

Solution

从所有度数为11的点出发一遍dfsdfs,对每一次遍历得到的字符串建入到SAMSAM。然后就是求本质不同的子串多少个。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#include<cstdio>
#include<algorithm>
#include<iostream>
#include<cstring>
#include<map>
#include<queue>
#include<bitset>
#define mk make_pair
#define fi first
#define nd second
#define pii pair<int,int>
#define pb push_back
#define sqr(x) ((x)*(x))
using namespace std;
typedef long long ll;
inline ll read() {ll x = 0; char ch = getchar(), w = 1;while(ch < '0' || ch > '9') {if(ch == '-') w = -1;
ch = getchar();}while(ch >= '0' && ch <= '9') {x = x * 10 + ch - '0';ch = getchar();}return x * w;}
void write(ll x) {if(x < 0) putchar('-'), x = -x;if(x > 9) write(x / 10);putchar(x % 10 + '0');}
inline void writeln(ll x) {write(x);puts("");}
const int N = 2e5 + 7;
int n;
struct SuffixAutoMaton {
int last;
int ch[N*10][11];
int num, len[N*10], fa[N*10], siz[N*10];
void insert(int c) {
int np = ++num, p = last; len[np] = len[last] + 1; last = np;
for(; p && !ch[p][c]; p = fa[p]) ch[p][c] = np;
if(!p) fa[np] = 1;
else {
int q = ch[p][c];
if(len[q] == len[p] + 1) fa[np] =q;
else {
int nq = ++num; memmove(ch[nq],ch[q],sizeof ch[q]);
len[nq] = len[p] + 1;
fa[nq] = fa[q];
fa[np] = fa[q] = nq;
for(; ch[p][c] == q; p = fa[p]) ch[p][c] = nq;
}
}
siz[last] = 1;
}
int c[N*10], a[N*10];
void init() {
ll ans = 0;
for(int i = 1; i <= num; ++i) ans += len[i] - len[fa[i]];
writeln(ans);
}
}sam;
int a[N];
int d[N];
struct Edge {
int u, v, nxt;
}e[N * 2];
int head[N], en;
void addedge(int x, int y) {
e[++en].u = x, e[en].v = y, e[en].nxt = head[x], head[x] = en;
}
void dfs(int x, int F) {
sam.insert(a[x]);
int tmp = sam.last;
for(int i = head[x]; i;i = e[i].nxt) {
int y = e[i].v;
if(y == F) continue;
sam.last = tmp;
dfs(y, x);
}
}
int main() {
sam.num = sam.last = 1;
n = read(); read();
for(int i = 1; i <= n; ++i) a[i] = read();
for(int i = 1; i < n; ++i) {
int x = read(), y = read();
addedge(x, y);
addedge(y, x);
++d[x], ++d[y];
}
for(int i = 1; i <= n; ++i) if(d[i] == 1) dfs(i, 0), sam.last = 1;
sam.init();
return 0;
}