廣義字尾自動機例題
阿新 • • 發佈:2021-08-11
以每個葉子節點為開頭進行dfs遍歷,將遍歷到的串全部加入建立廣義SAM,結果即為本質不同的字串個數
#include<cstdio> #include<cstring> #include<algorithm> #include<iostream> #include<queue> #define int long long using namespace std; const int maxn=2e6+10; int n,m,tot,len[maxn],fa[maxn],ch[maxn][26],c; int head[maxn<<1],nex[maxn<<1],to[maxn<<1]; int cnt[maxn]; char s[maxn]; int vis[maxn]; int res; int col[maxn]; int k=1; int ecnt; void add(int x,int y) { to[++ecnt]=y; nex[ecnt]=head[x]; head[x]=ecnt; } int Ins(int c,int last) { int p=last; if(ch[p][c]) { int q=ch[p][c]; if(len[p]+1==len[q]) { return q; } else { int nq=++tot; len[nq]=len[p]+1; memcpy(ch[nq],ch[q],sizeof(ch[nq])); fa[nq]=fa[q]; fa[q]=nq; for(; p&&ch[p][c]==q; p=fa[p])ch[p][c]=nq; return nq; } } int np=++tot; len[np]=len[p]+1; for(; p&&!ch[p][c]; p=fa[p])ch[p][c]=np; if(!p)fa[np]=1; else { int q=ch[p][c]; if(len[p]+1==len[q])fa[np]=q; else { int nq=++tot; vis[nq]=k; len[nq]=len[p]+1; memcpy(ch[nq],ch[q],sizeof(ch[nq])); fa[nq]=fa[q]; fa[q]=fa[np]=nq; for(; p&&ch[p][c]==q; p=fa[p])ch[p][c]=nq; } } return np; } int id[maxn]; inline void prepare() { for (int i=1; i<=tot; i++) cnt[len[i]]++; for (int i=1; i<=tot; i++) cnt[i]+=cnt[i-1]; for (int i=1; i<=tot; i++) id[cnt[len[i]]--]=i; for (int i=tot; i>=1; i--) { int X=id[i]; if (vis[fa[X]]==0) vis[fa[X]]=vis[X]; else if (vis[fa[X]]!=vis[X]) vis[fa[X]]=-1; } } void dfs(int x,int fa,int last) { int tmp=Ins(col[x],last); for(int i=head[x]; i; i=nex[i]) { int y=to[i]; if(y==fa)continue; dfs(y,x,tmp); } } int in[maxn]; signed main() { std::ios::sync_with_stdio(false); int n,c; cin>>n>>c; tot=1; for(int i=1; i<=n; i++) { cin>>col[i]; } for(int i=1; i<n; i++) { int x,y; cin>>x>>y; in[x]++; in[y]++; add(x,y); add(y,x); } for(int i=1; i<=n; i++) { if(in[i]==1) dfs(i,i,1); } for(int i=2; i<=tot; i++) { res+=len[i]-len[fa[i]]; } cout<<res<<endl; }