Luogu P3346 [ZJOI2015]诸神眷顾的幻想乡

广义 \(SAM\)

由于叶子节点比较少，直接暴力枚举叶子，把所有串插入广义 \(SAM\) 即可。

#include<iostream>
#include<cstdio>
#include<cstring>
#define ll long long
#define R register int
using namespace std;
namespace Luitaryi {
inline int g() { R x=0,f=1;
  register char s; while(!isdigit(s=getchar())) f=s=='-'?-1:f;
  do x=x*10+(s^48); while(isdigit(s=getchar())); return x*f;
} const int N=4000010;
int n;
int c[N];
ll ans;
int vr[N<<1],nxt[N<<1],fir[N],in[N],cnt;
inline void add(int u,int v) {
  vr[++cnt]=v,nxt[cnt]=fir[u],fir[u]=cnt;
  vr[++cnt]=u,nxt[cnt]=fir[v],fir[v]=cnt;
}
struct SAM {
  int tot,fa[N],c[N][10],len[N];
  SAM() {tot=1;}
  inline int add(int ch,int lst) {
    if(c[lst][ch]&&len[c[lst][ch]]==len[lst]+1)
      return c[lst][ch];
    R np=++tot,p=lst,q,nq,flg=0;
    len[np]=len[p]+1;
    while(p&&!c[p][ch]) c[p][ch]=np,p=fa[p];
    if(!p) fa[np]=1;
    else {
      q=c[p][ch];
      if(len[q]==len[p]+1) fa[np]=q;
      else {
        if(len[np]==len[p]+1) flg=1;
        nq=++tot;
        memcpy(c[nq],c[q],40);
        fa[nq]=fa[q],len[nq]=len[p]+1;
        fa[np]=fa[q]=nq;
        while(p&&c[p][ch]==q) c[p][ch]=nq,p=fa[p];
      }
    } return flg?nq:np;
  }
}s;
inline void dfs(int u,int fa,int lst) {
  R pos=s.add(c[u],lst);
  for(R i=fir[u];i;i=nxt[i]) {
    R v=vr[i];
    if(v==fa) continue;
    dfs(v,u,pos);
  }
}
inline void main() {
  n=g(),g();
  for(R i=1;i<=n;++i) c[i]=g();
  for(R i=1,u,v;i<n;++i) 
    u=g(),v=g(),add(u,v),++in[u],++in[v];
  for(R i=1;i<=n;++i) if(in[i]==1) dfs(i,0,1);
  for(R i=2;i<=s.tot;++i) 
    ans+=s.len[i]-s.len[s.fa[i]];
  //cout<<s.tot<<endl;
  printf("%lld\n",ans);
}
} signed main() {Luitaryi::main(); return 0;}

---

2020.01.10