[bzoj3926] [loj2137] [Zjoi2015] 诸神眷顾的幻想乡

Description

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的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$ 个。

Input

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

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

接下来 $n-1$ 行，每行两个正整数 $u$ , $v$ ，表示有一条连接空地 $u$ 和空地 $v$ 的边。

Output

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

Sample Input

7 3

0 2 1 2 1 0 0

1 2

3 4

3 5

4 6

5 7

2 5

Sample Output

HINT

对于所有数据，$1 \leq n \leq 100000$ , $1 \leq c \leq 10$。

对于 $15%$ 的数据，$n \leq 2000$。

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

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

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

想法

小数据其实已经提示做法了……

由题意，叶子节点不超过 $20$ 个。

经过手玩发现，分别以每个叶子节点为根，走到所有点，路径覆盖了所有点对之间的路径。

要求本质不同的子串数，这就是广义后缀自动机的模板题了。

啥叫广义后缀自动机？

就是多个字符串可以建在同一个后缀自动机上。

具体方法是，改 $last$ 的位置。

如插入完一个字符串，再插一个新串时，$last$ 改到 $root$

代码

调了半天，发现 $insert$ 里少写了一种情况 $orz...$

#include<cstdio>

#include<iostream>

#include<algorithm>

#include<cstring>

using namespace std;

int read(){

	char ch=getchar();

	int x=0;

	while(!isdigit(ch)) ch=getchar();

	while(isdigit(ch)) x=x*10+ch-'0',ch=getchar();

	return x;

}

const int N = 100005;

typedef long long ll;

struct node{

	int v;

	node *nxt;

}pool[N*2],*h[N];

int cnt;

void addedge(int u,int v){

	node *p=&pool[++cnt],*q=&pool[++cnt];

	p->v=v;p->nxt=h[u];h[u]=p;

	q->v=u;q->nxt=h[v];h[v]=q;

}

int n,C;

int col[N],deg[N];

struct trie{

	trie *ch[11],*pa;

	int len;

}pool2[N*40],*root,*pos[N],*last;

int cnt2;

void insert(int c){

	trie *p=last,*cur=&pool2[++cnt2];

	cur->len=p->len+1;

	for(;p && !p->ch[c];p=p->pa) p->ch[c]=cur;

	if(!p) cur->pa=root;

	else {

		trie *q=p->ch[c],*nq;

		if(q->len!=p->len+1){

			nq=&pool2[++cnt2];

			for(int i=0;i<C;i++) nq->ch[i]=q->ch[i];

			nq->len=p->len+1;

			nq->pa=q->pa;

			q->pa=cur->pa=nq;

			for(;p && p->ch[c]==q;p=p->pa) p->ch[c]=nq;

		}

		else cur->pa=q;

	}

	last=cur;

}

int vis[N];

void dfs(int u){

	int v;

	vis[u]=1;

	for(node *p=h[u];p;p=p->nxt){

		if(vis[v=p->v]) continue;

		last=pos[u];

		insert(col[v]);

		pos[v]=last;

		dfs(v);

	}

}

ll dp[N*40];

ll count(trie *p){

	if(dp[p-pool2]!=-1) return dp[p-pool2];

	dp[p-pool2]=1;

	for(int i=0;i<C;i++)

		if(p->ch[i]) dp[p-pool2]+=count(p->ch[i]);

	return dp[p-pool2];

}

int main()

{

	int u,v;

	n=read(); C=read();

	for(int i=1;i<=n;i++) col[i]=read();

	for(int i=1;i<n;i++){

		u=read(); v=read();

		addedge(u,v);

		deg[u]++; deg[v]++;

	}

	root=&pool2[++cnt2];

	for(int i=1;i<=n;i++){

		if(deg[i]!=1) continue;

		memset(vis,0,sizeof(vis));

		last=root;

		insert(col[i]);

		pos[i]=last;

		dfs(i);

	}

	memset(dp,-1,sizeof(dp));

	printf("%lld\n",count(root)-1);

	return 0;

}