幽香是全幻想乡里最受人欢迎的萌妹子,这天,是幽香的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的边。
对于所有数据,1<=n<=100000, 1<=c<=10。
对于15%的数据,n<=2000。
另有5%的数据,所有空地都至多与两个空地相邻。
另有5%的数据,除一块空地与三个空地相邻外,其他空地都分别至多与两个空地相邻。
另有5%的数据,除某两块空地与三个空地相邻外,其他空地都分别至多与两个空地相邻
1 #include <bits/stdc++.h>
2
3 /* MYSCANNER */
4
5 #define siz 1024
6
7 inline int get_c(void)
8 {
9 static char buf[siz];
10 static char *head = buf + siz;
11 static char *tail = buf + siz;
12
13 if (head == tail)
14 fread(head = buf, 1, siz, stdin);
15
16 return *head++;
17 }
18
19 inline int get_i(void)
20 {
21 register int ret = 0;
22 register int neg = false;
23 register int bit = get_c();
24
25 for (; bit < 48; bit = get_c())
26 if (bit == ‘-‘)neg ^= true;
27
28 for (; bit > 47; bit = get_c())
29 ret = ret * 10 + bit - ‘0‘;
30
31 return neg ? -ret : ret;
32 }
33
34 const int maxn = 2000100;
35
36 /* AUTOMATON */
37
38 int tail = 1;
39 int root[maxn];
40 int step[maxn];
41 int fail[maxn];
42 int next[maxn][15];
43
44 inline int extend(int p, int c)
45 {
46 if (next[p][c] && step[next[p][c]] == step[p] + 1)
47 return next[p][c];
48 int t = tail++;
49 step[t] = step[p] + 1;
50 while (p && !next[p][c])
51 next[p][c] = t, p = fail[p];
52 if (p)
53 {
54 int q = next[p][c];
55 if (step[q] == step[p] + 1)
56 fail[t] = q;
57 else
58 {
59 int k = tail++;
60 fail[k] = fail[q];
61 fail[q] = fail[t] = k;
62 step[k] = step[p] + 1;
63 memcpy(next[k], next[q], sizeof(next[k]));
64 while (p && next[p][c] == q)
65 next[p][c] = k, p = fail[p];
66 }
67 }
68 else
69 fail[t] = 1;
70 return t;
71 }
72
73 /* GRAPH */
74
75 int n;
76 int m;
77 int tot;
78 int hd[maxn];
79 int nt[maxn];
80 int to[maxn];
81 int cnt[maxn];
82 int col[maxn];
83 int vis[maxn];
84 int que[maxn];
85
86 inline void addEdge(int x, int y)
87 {
88 nt[++tot] = hd[x]; to[tot] = y; hd[x] = tot;
89 nt[++tot] = hd[y]; to[tot] = x; hd[y] = tot;
90 }
91
92 void DFS(int u, int f)
93 {
94 root[u] = extend(root[f], col[u]);
95 for (int i = hd[u]; i; i = nt[i])
96 if (to[i] != f)DFS(to[i], u);
97 }
98
99 /* MAIN FUNC */
100
101 signed main(void)
102 {
103 n = get_i();
104 m = get_i();
105
106 for (int i = 1; i <= n; ++i)
107 col[i] = get_i();
108
109 for (int i = 1; i < n; ++i)
110 {
111 int x = get_i();
112 int y = get_i();
113 addEdge(x, y);
114 ++cnt[x];
115 ++cnt[y];
116 }
117
118 root[0] = tail++;
119
120 for (int i = 1; i <= n; ++i)
121 if (cnt[i] == 1)DFS(i, 0);
122
123 long long ans = 0LL;
124
125 int head = 0, tail = 0;
126
127 que[tail++] = 1;
128
129 while (head != tail)
130 {
131 int t = que[head++];
132 if (t != 1)
133 ans += step[t] - step[fail[t]];
134 for (int i = 0; i < m; ++i)
135 if (next[t][i] && !vis[next[t][i]])
136 vis[next[t][i]] = 1, que[tail++] = next[t][i];
137 }
138
139 printf("%lld\n", ans);
140 }