1. 程式人生 > 實用技巧 >[ZJOI2015]諸神眷顧的幻想鄉

[ZJOI2015]諸神眷顧的幻想鄉

[ZJOI2015]諸神眷顧的幻想鄉

題意

樹上每個點代表一個字元,問樹上路徑可以組成多少不同的字串(葉子最多有20個)

題解

樹上的路徑一定是葉子為根的一條從上到下的路徑,而葉子有很小,我們就可以把所有葉子構成的 trie 建成一個廣義的 sam,最後統計不同字串個數即可

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int N=4000055; 
int tot,to[200005],co[100005],nextt[200005],head[100005],n,c,de[100005]; 
int cnt,last,l[N],ch[N][10],fa[N];
long long ans;
inline int read()
{
	int k=0,f=1;char c=getchar();
	while(!isdigit(c)) {if(c=='-')f=-1;c=getchar();}
	while(isdigit(c)) k=k*10+c-'0',c=getchar(); return f*k;
}
void add(int a,int b)
{
	to[++tot]=b;
	nextt[tot]=head[a];
	head[a]=tot;
}
int ins(int c,int p)
{
	if(ch[p][c])
	{
		int q=ch[p][c];
		if(l[p]+1==l[q]) last=q;
		else 
		{
			int nq=++cnt;l[nq]=l[p]+1;last=nq;
			memcpy(ch[nq],ch[q],sizeof(ch[q]));
			fa[nq]=fa[q];fa[q]=nq;
			for(;ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
		}
	}
	else 
	{
		int np=++cnt;last=np;l[np]=l[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(l[p]+1==l[q]) fa[np]=q;
			else 
			{
				int nq=++cnt;l[nq]=l[p]+1;
				memcpy(ch[nq],ch[q],sizeof(ch[q]));
				fa[nq]=fa[q];fa[np]=fa[q]=nq;
				for(;ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
			}
		 } 
	}
	return last;
 } 

void dfs(int u,int f,int p)
{
	int t=ins(co[u],p);
	for(int i=head[u];i;i=nextt[i])
	{
		if(to[i]!=f)
		{
			dfs(to[i],u,t);
		}
	}
}
int main()
{
	int a,b;
	n=read();c=read();cnt=last=1;
	for(int i=1;i<=n;i++) co[i]=read();
	for(int i=1;i<n;i++) a=read(),b=read(),add(a,b),add(b,a),de[a]++,de[b]++;
	for(int i=1;i<=n;i++)
	{
		if(de[i]==1)
		dfs(i,0,1);
	}
	for(int i=1;i<=cnt;i++) ans+=l[i]-l[fa[i]];
	printf("%lld\n",ans);
	return 0;
}