[BZOJ3926] [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

30

Solution

注意到叶子节点只有\(20\)个，那么可以很暴力的以每个叶子节点为根，把树插入\(trie\)树中，那么任意一条路径都会是这个\(trie\)树的子串。

所以对这个\(trie\)树建立广义后缀自动机，那么答案就是：

\[\sum_{i=2}^{cnt}maxl_i-maxl_{par_i} \]

写代码的时候直接\(dfs\)就好了，没必要把\(trie\)建出来。

#include<bits/stdc++.h>
using namespace std;

#define ll long long

void read(int &x) {
    x=0;int f=1;char ch=getchar();
    for(;!isdigit(ch);ch=getchar()) if(ch=='-') f=-f;
    for(;isdigit(ch);ch=getchar()) x=x*10+ch-'0';x*=f;
}
 
void print(ll x) {
    if(x<0) putchar('-'),x=-x;
    if(!x) return ;print(x/10),putchar(x%10+48);
}
void write(ll x) {if(!x) putchar('0');else print(x);putchar('\n');}

const int maxn = 2e6+10;

struct edge{int to,nxt;}e[maxn<<1];
int a[maxn],head[maxn],tot,deg[maxn],n,_;

void ins(int u,int v) {e[++tot]=(edge){v,head[u]},head[u]=tot;}

namespace sam {   //Suffix Automaton 
	int par[maxn],ml[maxn],tr[maxn][10],cnt=1,qs=1;

	int append(int x,int lstp) {
		int p=lstp,q=tr[p][x];
		if(q) {
			if(ml[p]+1<ml[q]) {
				int nq=++cnt;ml[nq]=ml[p]+1;lstp=nq;
				memcpy(tr[nq],tr[q],sizeof tr[nq]);
				par[nq]=par[q],par[q]=nq;
				for(;p&&tr[p][x]==q;p=par[p]) tr[p][x]=nq;
			} else lstp=q;
		} else {
			int np=++cnt;ml[np]=ml[p]+1,lstp=np;
			for(;p&&tr[p][x]==0;p=par[p]) tr[p][x]=np;
			if(!p) par[np]=qs;
			else {
				q=tr[p][x];
				if(ml[p]+1<ml[q]) {
					int nq=++cnt;ml[nq]=ml[p]+1;
					memcpy(tr[nq],tr[q],sizeof tr[nq]);
					par[nq]=par[q],par[q]=par[np]=nq;
					for(;p&&tr[p][x]==q;p=par[p]) tr[p][x]=nq;
				} else par[np]=q;
			}
		}
		return lstp;
	}

	void dfs(int x,int fa,int lstp) {
		lstp=append(a[x],lstp);
		for(int i=head[x];i;i=e[i].nxt)
			if(e[i].to!=fa) dfs(e[i].to,x,lstp);
	}
	
	ll solve() {
		ll res=0;
		for(int i=2;i<=cnt;i++) res+=ml[i]-ml[par[i]];
		return res;
	}
}

int main() {
	read(n);read(_);for(int i=1;i<=n;i++) read(a[i]);
	for(int i=1,x,y;i<n;i++) read(x),read(y),ins(x,y),ins(y,x),deg[x]++,deg[y]++;
	for(int i=1;i<=n;i++) if(deg[i]==1) sam :: dfs(i,0,1);
	write(sam :: solve());
	return 0;
}