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;
}

bzoj 3926

原文：http://www.cnblogs.com/zqq123/p/5119721.html