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

bzoj
luogu

题意

有一棵树，每个节点上有一个字符，一条路径上的字符连起来就是一个字符串。求树上一共有多少个不同的字符串。
\(n\le10^5\)，字符集大小\(c\le10\)，只与一个空地相邻的空地数量不超过20个。

sol

最后一句话是这道题的关键。
什么叫做“只与一个空地相邻的空地数量不超过20个”？
叶子节点数量不超过20个！
我们只要从每个叶子节点出发dfs整棵树，就可以得到所有的子串了。
所以当然是一边dfs一遍extend，注意广义后缀自动机，要保存当前节点的前一个点\(last\)的值。每次extend前要务必保证这点。
不同的子串个数？\(\sum_{i=1}^{tot}len[i]-len[fa[i]]\)！

code

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
int gi()
{
	int x=0,w=1;char ch=getchar();
	while ((ch<'0'||ch>'9')&&ch!='-') ch=getchar();
	if (ch=='-') w=0,ch=getchar();
	while (ch>='0'&&ch<='9') x=(x<<3)+(x<<1)+ch-'0',ch=getchar();
	return w?x:-x;
}
const int N = 4e6+5;
int n,col[N],to[N],nxt[N],head[N],cnt,d[N],tr[N][10],fa[N],len[N],last=1,tot=1;
void link(int u,int v){to[++cnt]=v;nxt[cnt]=head[u];head[u]=cnt;}
void extend(int c)
{
	int v=last,u=++tot;last=u;
	len[u]=len[v]+1;
	while (v&&!tr[v][c]) tr[v][c]=u,v=fa[v];
	if (!v) fa[u]=1;
	else{
		int x=tr[v][c];
		if (len[x]==len[v]+1) fa[u]=x;
		else{
			int y=++tot;
			memcpy(tr[y],tr[x],sizeof(tr[y]));
			fa[y]=fa[x];fa[x]=fa[u]=y;len[y]=len[v]+1;
			while (v&&tr[v][c]==x) tr[v][c]=y,v=fa[v];
		}
	}
}
void dfs(int u,int f)
{
	extend(col[u]);int tmp=last;
	for (int e=head[u];e;e=nxt[e])
		if (to[e]!=f)
			last=tmp,dfs(to[e],u);
}
int main()
{
	n=gi();gi();long long ans=0;
	for (int i=1;i<=n;++i) col[i]=gi();
	for (int i=1;i<n;++i)
	{
		int u=gi(),v=gi();
		link(u,v);link(v,u);++d[u];++d[v];
	}
	for (int i=1;i<=n;++i) if (d[i]==1) last=1,dfs(i,0);
	for (int i=1;i<=tot;++i) ans+=len[i]-len[fa[i]];
	printf("%lld\n",ans);return 0;
}