题目描述
幽香是全幻想乡里最受人欢迎的萌妹子,这天,是幽香的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:
7 3
0 2 1 2 1 0 0
1 2
3 4
3 5
4 6
5 7
2 5
输出样例#1:
30
说明
对于所有数据,1<=n<=100000, 1<=c<=10。
对于15%的数据,n<=2000。
另有5%的数据,所有空地都至多与两个空地相邻。
另有5%的数据,除一块空地与三个空地相邻外,其他空地都分别至多与两个空地相邻。
另有5%的数据,除某两块空地与三个空地相邻外,其他空地都分别至多与两个空地相邻
分析:
不同子串个数显然可以用后缀自动机,答案就是所有点的
的大小。
因为叶子节点不超过
个,以这
个节点为根的trie建广义后缀自动机。
至于建广义后缀自动机,记录一个点父亲节点对于的字符在trie上的位置,记为
。当加入一个字符时,如果这个串还没有出现过,就同普通的自动机一样处理;如果出现过这个串,那么相当于要不要新建一个点那个部分那样讨论,因为可能出现
集的改变。
代码:
#include <iostream>
#include <cstdio>
#include <cmath>
#define LL long long
const int maxn=1e5+7;
const int maxp=2e6+7;
using namespace std;
int n,c,x,y,cnt,last;
int ls[maxn],a[maxn],ru[maxn];
LL ans;
struct node{
int len,fail;
int son[12];
}t[maxp];
struct edge{
int y,next;
}g[maxn*2];
void add(int x,int y)
{
g[++cnt]=(edge){y,ls[x]};
ls[x]=cnt;
}
int ins_sam(int x,int c)
{
int now,p,q,clone;
p=x;
if (t[p].son[c])
{
q=t[p].son[c];
if (t[p].len==t[q].len) return q;
clone=++cnt;
t[clone]=t[q];
t[clone].len=t[p].len+1;
t[q].fail=clone;
while (p&&(t[p].son[c]==q)) t[p].son[c]=clone,p=t[p].fail;
return clone;
}
now=++cnt;
t[now].len=t[p].len+1;
while (p&&(!t[p].son[c])) t[p].son[c]=now,p=t[p].fail;
if (!p) t[now].fail=1;
else
{
q=t[p].son[c];
if (t[p].len+1==t[q].len) t[now].fail=q;
else
{
clone=++cnt;
t[clone]=t[q];
t[clone].len=t[p].len+1;
t[now].fail=t[q].fail=clone;
while (p&&(t[p].son[c]==q)) t[p].son[c]=clone,p=t[p].fail;
}
}
return now;
}
void dfs(int x,int fa,int last)
{
int d=ins_sam(last,a[x]);
for (int i=ls[x];i>0;i=g[i].next)
{
int y=g[i].y;
if (y==fa) continue;
dfs(y,x,d);
}
}
int main()
{
scanf("%d%d",&n,&c);
for (int i=1;i<=n;i++) scanf("%d",&a[i]);
for (int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
add(x,y);
add(y,x);
ru[x]++,ru[y]++;
}
cnt=1;
for (int i=1;i<=n;i++)
{
if (ru[i]==1) dfs(i,0,1);
}
for (int i=2;i<=cnt;i++)
ans+=(LL)t[i].len-(LL)t[t[i].fail].len;
printf("%lld",ans);
}