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
HINT
对于所有数据，1<=n<=100000, 1<=c<=10。
对于15%的数据，n<=2000。
另有5%的数据，所有空地都至多与两个空地相邻。
另有5%的数据，除一块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。
另有5%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻

---

思路

因为叶子节点很少，所以对于每个叶子节点直接遍历整个树
然后把路径上的串全部插进一个trie树中
然后对这个trie建广义后缀自动机
注意这里的last不能是全局的last，应该是trie上父亲节点的sam节点
然后直接统计不同串的个数就好了

---

#include<bits/stdc++.h>
using namespace std;
#define fu(a,b,c) for(int a=b;a<=c;++a)
#define fd(a,b,c) for(int a=b;a>=c;--a)
#define N 100001
#define LL long long
const int CHARSET_SIZE=10;
struct Edge{int v,next;}E[N<<1];
int head[N],tot=0,deg[N]={0};
void add(int u,int v){
  E[++tot]=(Edge){v,head[u]};
  head[u]=tot;
}
int n,c,col[N];
//Suffix_Automaton
struct Sam{
  Sam *ch[CHARSET_SIZE],*prt;
  int maxl;
  Sam(int maxl=0):maxl(maxl){}
};
struct Suffix_Automaton{
  Sam pool[N*40],*cur,*root;
  Suffix_Automaton(){
    cur=pool;
    root=new (cur++)Sam;
  }
  Sam *extend(int c,Sam *last){
    Sam *u=new (cur++)Sam(last->maxl+1),*v=last;
    for(;v&&!v->ch[c];v=v->prt)v->ch[c]=u;
    if(!v)u->prt=root;
    else if(v->ch[c]->maxl==v->maxl+1){
      u->prt=v->ch[c];
    }else{
      Sam *n=new (cur++)Sam(v->maxl+1),*o=v->ch[c];
      copy(o->ch,o->ch+CHARSET_SIZE,n->ch);
      n->prt=o->prt;
      o->prt=u->prt=n;
      for(;v&&v->ch[c]==o;v=v->prt)v->ch[c]=n;
    }
    return u;
  }
  LL solve(){
    LL ans=0;
    for(Sam *p=pool+1;p!=cur;p++)
      ans+=p->maxl-p->prt->maxl;
    return ans;
  }
}sam;
//Trie
struct Node{
  Node *ch[CHARSET_SIZE];
};
struct Trie{
  Node pool[N*25],*cur,*root,*t[N];
  Trie(){
    cur=pool;
    t[0]=root=new (cur++)Node;
  }
  Node *insert(Node *last,int c){
    if(!last->ch[c])last->ch[c]=new (cur++)Node;
    return last->ch[c];
  }
  #define pi pair<Node*,Sam*>
  void bfs(){
    queue<pi> q;
    q.push(pi(root,sam.root));
    while(!q.empty()){
      pi u=q.front();q.pop();
      fu(i,0,c-1)
        if(u.first->ch[i]){
          Sam *tmp=sam.extend(i,u.second);
          q.push(pi(u.first->ch[i],tmp));
        }
    }
  }
}trie;
//Main and Solve
void dfs(int u,int fa){
  trie.t[u]=trie.insert(trie.t[fa],col[u]);
  for(int i=head[u];i;i=E[i].next){
    int v=E[i].v;
    if(v==fa)continue;
    dfs(v,u);
  }
}
int main(){
  scanf("%d%d",&n,&c);
  fu(i,1,n)scanf("%d",&col[i]);
  fu(i,2,n){
    int u,v;scanf("%d%d",&u,&v);
    add(u,v);
    add(v,u);
    deg[u]++;deg[v]++;
  }
  fu(i,1,n)if(deg[i]==1)dfs(i,0);
  trie.bfs();
  printf("%lld",sam.solve());
  return 0;
}