BZOJ3926 ZJOI2015 诸神眷顾的幻想乡 Trie、广义SAM

传送门

---

树上的任意一条路径一定会在以某一个叶子节点为根的树上成为一条直上直下的链，而总共只有\(20\)个叶子节点。

于是每一次选所有叶子节点中的一个作为根，形成一个\(Trie\)，把\(20\)个\(Trie\)统一到一个\(Trie\)上，然后对这个总的\(Trie\)建立广义后缀自动机，最后统计一下广义SAM每个节点代表的字符串个数求和。

在\(Trie\)上BFS建立广义SAM可以做到\(O(20NA)\)的复杂度，把\(20\)个\(Trie\)直接一个个插入SAM在线建会多一个\(O(\text{总Trie上所有叶子节点的深度和})\)，证明见2015lyy的论文。

#include<iostream>
#include<cstdio>
#include<cctype>
#include<algorithm>
#include<cstring>
#include<queue>
//This code is written by Itst
using namespace std;

inline int read(){
    int a = 0;
    char c = getchar();
    while(!isdigit(c))
        c = getchar();
    while(isdigit(c)){
        a = a * 10 + c - 48;
        c = getchar();
    }
    return a;
}

const int MAXN = 4e6 + 7 , MAXM = 1e5 + 7;
namespace SAM{
    int trans[MAXN][10] , fa[MAXN] , Lst[MAXN] , Sst[MAXN];
    int N , cnt = 1;

    int insert(int p , int l , int x){
        int t = ++cnt;
        Lst[t] = l;
        while(p && !trans[p][x]){
            trans[p][x] = t;
            p = fa[p];
        }
        if(!p){
            fa[t] = Sst[t] = 1;
            return t;
        }
        int q = trans[p][x];
        Sst[t] = Lst[p] + 2;
        if(Lst[q] == Lst[p] + 1){
            fa[t] = q;
            return t;
        }
        int k = ++cnt;
        memcpy(trans[k] , trans[q] , sizeof(trans[q]));
        Lst[k] = Lst[p] + 1; Sst[k] = Sst[q];
        Sst[q] = Lst[p] + 2;
        fa[k] = fa[q]; fa[q] = fa[t] = k;
        while(trans[p][x] == q){
            trans[p][x] = k;
            p = fa[p];
        }
        return t;
    }

    long long calc(){
        long long ans = 0;
        for(int i = 2 ; i <= cnt ; ++i)
            ans += Lst[i] - Sst[i] + 1;
        return ans;
    }
}
struct Edge{
    int end , upEd;
}Ed[MAXM << 1];
int col[MAXM] , head[MAXM] , in[MAXM] , ch[MAXN][10] , ind[MAXN] , dep[MAXN];
int N , C , cnt = 1 , cntEd;

inline void addEd(int a , int b){
    Ed[++cntEd] = Edge{b , head[a]};
    head[a] = cntEd;
    ++in[b];
}

void dfs(int x , int p , int cur){
    if(!ch[cur][col[x]])
        ch[cur][col[x]] = ++cnt;
    cur = ch[cur][col[x]];
    for(int i = head[x] ; i ; i = Ed[i].upEd)
        if(Ed[i].end != p)
            dfs(Ed[i].end , x , cur);
}

void create(){
    ind[1] = 1;
    queue < int > q;
    q.push(1);
    while(!q.empty()){
        int t = q.front();
        q.pop();
        for(int i = 0 ; i < 10 ; ++i)
            if(ch[t][i]){
                dep[ch[t][i]] = dep[t] + 1;
                ind[ch[t][i]] = SAM::insert(ind[t] , dep[ch[t][i]] , i);
                q.push(ch[t][i]);
            }
    }
}

int main(){
#ifndef ONLINE_JUDGE
    freopen("in","r",stdin);
    //freopen("out","w",stdout);
#endif
    N = read();
    C = read();
    for(int i = 1 ; i <= N ; ++i)
        col[i] = read();
    for(int i = 1 ; i < N ; ++i){
        int a = read() , b = read();
        addEd(a , b); addEd(b , a);
    }
    for(int i = 1 ; i <= N ; ++i)
        if(in[i] == 1)
            dfs(i , 0 , 1);
    create();
    cout << SAM::calc();
    return 0;
}