字符串（广义后缀自动机）：BZOJ 3926 [Zjoi2015]诸神眷顾的幻想乡

 1 #include <iostream>
 2 #include <cstring>
 3 #include <cstdio>
 4 using namespace std;
 5 const int maxn=100010;
 6 const int maxm=2000010;
 7 long long ans;
 8 bool vis[maxm];
 9 int n,sumc,cntE,c[maxn];
10 int fir[maxn],nxt[maxn<<1],to[maxn<<1];
11 void addedge(int a,int b){
12     nxt[++cntE]=fir[a];
13     fir[a]=cntE;
14     to[cntE]=b;
15 }
16 int st[maxm],topst;
17 int CH[maxm][11],pos[maxm];
18 int fa[maxm],ch[maxm][11],len[maxm];
19 struct SAM{
20     int cnt,last;
21     SAM(){
22         cnt=last=1;
23     }
24     int Insert(int c){
25         int p=last,np=last=++cnt;len[np]=len[p]+1;
26         while(p&&!ch[p][c])ch[p][c]=np,p=fa[p];
27         if(!p)fa[np]=1;
28         else{
29             int q=ch[p][c];
30             if(len[p]+1==len[q])
31                 fa[np]=q;
32             else{
33                 int nq=++cnt;len[nq]=len[p]+1;
34                 memcpy(ch[nq],ch[q],sizeof(ch[q]));
35                 fa[nq]=fa[q];fa[q]=fa[np]=nq;
36                 while(ch[p][c]==q)ch[p][c]=nq,p=fa[p];
37             }    
38         }
39         ans+=len[last]-len[fa[last]];
40         return last;
41     }
42 }sam;
43 struct Trie{
44     int cnt,rt;
45     Trie(){
46         rt=cnt=1;
47         pos[rt]=1;
48     }
49     void Insert(int &x,int p){
50         vis[p]=true;
51         if(!x)x=++cnt;
52         for(int i=fir[p];i;i=nxt[i])
53             if(!vis[to[i]])
54                 Insert(CH[x][c[to[i]]],to[i]);
55     }
56     void Construct(int x){
57         for(int i=0;i<sumc;i++)
58             if(CH[x][i]){
59                 sam.last=pos[x];
60                 pos[CH[x][i]]=sam.Insert(i);
61             }
62         for(int i=0;i<sumc;i++)
63             if(CH[x][i])
64                 Construct(CH[x][i]);
65     }
66 }trie;
67 int main(){
68 #ifndef ONLINE_JUDGE
69     freopen("zjoi15_substring.in","r",stdin);
70     freopen("zjoi15_substring.out","w",stdout);
71 #endif
72     scanf("%d%d",&n,&sumc);
73     for(int i=1;i<=n;i++)
74         scanf("%d",&c[i]); 
75     for(int i=1,a,b;i<n;i++){
76         scanf("%d%d",&a,&b);
77         addedge(a,b);
78         addedge(b,a);
79     }    
80     for(int x=1,tot=0;x<=n;tot=0,x++){
81         for(int i=fir[x];i;i=nxt[i])tot+=1;
82         if(tot==1)st[++topst]=x;
83     }
84     while(topst){
85         memset(vis,0,sizeof(vis));
86         trie.Insert(CH[trie.rt][c[st[topst]]],st[topst]);
87         topst-=1;
88     }
89     trie.Construct(trie.rt);
90     printf("%lld\n",ans);
91     return 0;    
92 }