bzoj 3926

bzoj 3926

trie树后缀自动机，直接用val值统计答案

注意细节，p，q不要打反；注意删调

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define maxn 2000020
#define maxm 100020

struct node{
	int next[12];
}trie[maxn];
struct node1{
	int next[12],val,pnt;
}sam[maxn * 2];
struct node2{
	int to,next;
}e[maxm * 2];
int tot,cnt,head[maxm],root;
long long sum;
int n,col[maxm],c,degree[maxm],last;

inline void addedge(int x,int y){
	e[++cnt].to = y;
	e[cnt].next = head[x];
	head[x] = cnt;
	degree[x]++;
}
inline int add_sam(int x,int last){
 int np = ++tot;
 int p = last;
 sam[np].val = sam[p].val + 1;
 while ( p && !sam[p].next[x] ){
 sam[p].next[x] = np , p = sam[p].pnt;
 }
 int q = sam[p].next[x];
 if ( !q ){
 sam[p].next[x] = np;
 sam[np].pnt = p;
 } 
 else if ( q && sam[p].val + 1 == sam[q].val ){
 sam[np].pnt = q;
 }
 else{
 int nq = ++tot;
 sam[nq].pnt = sam[q].pnt;
 sam[np].pnt = sam[q].pnt = nq;
 sam[nq].val = sam[p].val + 1;
 memcpy(sam[nq].next,sam[q].next,sizeof(sam[q].next));
 while ( p && sam[p].next[x] == q ){
 sam[p].next[x] = nq , p = sam[p].pnt;
 }
 if ( sam[p].next[x] == q ) sam[p].next[x] = nq;
 }
 return np;
}
void dfs_trie(int now,int last){
	int cur = last;
	for (int i = 0 ; i <= c; i++){
	if ( trie[now].next[i] ){
	last = add_sam(i,cur);
	dfs_trie(trie[now].next[i],last);
	}
	}
}
void dfs(int now,int &trie_now,int fa){
	if ( !trie_now ) trie_now = ++tot;
	for (int i = head[now] ; i ; i = e[i].next){
	if ( fa == e[i].to ) continue;
	dfs(e[i].to,trie[trie_now].next[col[e[i].to]],now);
	} 
}
int main(){
	scanf("%d %d",&n,&c);
	for (int i = 1 ; i <= n ; i++) scanf("%d",&col[i]);
	for (int i = 1 ; i < n ; i++){
	int u,v;
	scanf("%d %d",&u,&v);
	addedge(u,v);
	addedge(v,u);
	}
	for (int i = 1 ; i <= n ; i++){
	if ( degree[i] == 1 ){
	dfs(i,trie[0].next[col[i]],0);
	}
	}
	tot = 0;
	dfs_trie(0,root);
	for (int i = 1 ; i <= tot ; i++){
	sum += (long long) sam[i].val - (long long) sam[sam[i].pnt].val;
	}
	printf("%lld",sum);
	return 0;
}