BZOJ3926 [Zjoi2015]诸神眷顾的幻想乡【广义后缀自动机】

//#pragma GCC optimize("O3")
//#pragma comment(linker, "/STACK:1024000000,1024000000")
#include<bits/stdc++.h>
using namespace std;
function<void(void)> ____ = [](){ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);};
const int MAXN = 4e6+7;
int n,m,col[MAXN];
vector<int> G[MAXN];
struct EXSAM{
    int len[MAXN],link[MAXN],ch[MAXN][10],tot,last,sta[MAXN],par[MAXN];
    EXSAM(){ tot = last = 1; }
    void extend(int c){
        int np = 0, p = last;
        if(!ch[p][c]) last = np = ++tot, len[np] = len[p] + 1;
        while(p and !ch[p][c]){
            ch[p][c] = np;
            p = link[p];
        }
        if(!p) link[np] = 1;
        else{
            int q = ch[p][c];
            if(len[p] + 1 == len[q]) np?link[np]=q:last=q;
            else{
                int clone = ++tot;
                len[clone] = len[p] + 1;
                link[clone] = link[q];
                memcpy(ch[clone],ch[q],sizeof(ch[q]));
                while(p and ch[p][c]==q){
                    ch[p][c] = clone;
                    p = link[p];
                }
                link[q] = clone;
                if(np) link[np] = clone;
                else last = clone;
            }
        }
    }
    void bfs(int u){
        queue<int> que;
        que.push(u);
        sta[par[u]=0] = 1;
        while(!que.empty()){
            u = que.front(); que.pop();
            last = sta[par[u]], extend(col[u]), sta[u] = last;
            for(int i = 0; i < (int)G[u].size(); i++) if(G[u][i]!=par[u]){
                par[G[u][i]] = u;
                que.push(G[u][i]);
            }       
        }
    }
    long long query(){
        long long ret = 0;
        for(int i = 2; i <= tot; i++) ret += len[i] - len[link[i]];
        return ret;
    }
}sam;
int main(){
    scanf("%d %d",&n,&m);
    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);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    for(int i = 1; i <= n; i++) if(G[i].size()==1) sam.bfs(i);
    printf("%lld\n",sam.query());
    return 0;
}

BZOJ3926 [Zjoi2015]诸神眷顾的幻想乡【广义后缀自动机】

原文：https://www.cnblogs.com/kikokiko/p/12718180.html