Luogu3346[ZJOI2015]诸神眷顾的幻想乡

Luogu3346[ZJOI2015]诸神眷顾的幻想乡

题面：洛谷

解析

观察到树的叶子节点数量很少（题面描述是真的迷，我开始以为是度数小于20），发现对于树上的每一条路径，都能在以某一个叶子节点为根的树中表示为一条从上到下的路径，那么把每一颗树看做一颗\(Trie\)树，用广义后缀自动机插入即可，答案就是不同状态与他后缀链接的状态\(len\)之差的和。

代码

// luogu-judger-enable-o2
#include<cstdio>
#include<cstring>
#define N 100005
#define LL long long
using namespace std;
int n,c,col[N];
inline int In(){
    char c=getchar(); int x=0,ft=1;
    for(;c<'0'||c>'9';c=getchar()) if(c=='-') ft=-1;
    for(;c>='0'&&c<='9';c=getchar()) x=x*10+c-'0';
    return x*ft;
}
namespace Sam{
    int root,sz,ch[N*20][10],fa[N*20],len[N*20];
    inline void init(){
        sz=0; root=++sz;
    }
    inline int newnode(int u){
        len[++sz]=len[u]+1; return sz;
    }
    inline int Extend(int f,int c){
        if(ch[f][c]&&len[ch[f][c]]==len[f]+1) return ch[f][c];
        int p=newnode(f);
        while(f&&!ch[f][c]) ch[f][c]=p,f=fa[f];
        if(!f){ fa[p]=root; return p; }
        int x=ch[f][c],ff=0;
        if(len[x]==len[f]+1){ fa[p]=x; return p; }
        if(len[p]==len[f]+1) ff=1;
        int y=newnode(f); memcpy(ch[y],ch[x],sizeof(ch[y]));
        fa[y]=fa[x]; fa[x]=fa[p]=y;
        while(f&&ch[f][c]==x) ch[f][c]=y,f=fa[f];
        return ff?y:p;
    }
    inline void Calc(){
        LL ans=0;
        for(int i=root;i<=sz;++i) ans+=len[i]-len[fa[i]];
        printf("%lld\n",ans);
    }
}
int h[N],deg[N],e_tot=0;
struct E{ int to,nex; }e[N<<1];
inline void add(int u,int v){
    e[++e_tot]=(E){v,h[u]}; h[u]=e_tot; ++deg[u];
    e[++e_tot]=(E){u,h[v]}; h[v]=e_tot; ++deg[v];
}
void dfs(int u,int las,int fa){
    int v_las=Sam::Extend(las,col[u]);
    for(int i=h[u],v;i;i=e[i].nex){
        v=e[i].to; if(v==fa) continue;
        dfs(v,v_las,u);
    }
}
int main(){
    n=In(); c=In(); Sam::init();
    for(int i=1;i<=n;++i) col[i]=In();
    for(int i=1;i<n;++i) add(In(),In());
    for(int i=1;i<=n;++i) if(deg[i]==1) dfs(i,Sam::root,-1);
    Sam::Calc();
    return 0;
}