BZOJ3926 [Zjoi2015]诸神眷顾的幻想乡 字符串 SAM

原文链接https://www.cnblogs.com/zhouzhendong/p/BZOJ3926.html

题目传送门 - BZOJ3926

题意

　　给定一个有 $n$ 个节点，最多只有 $20$ 个度为 $1$ 的节点的树。

　　树上每一个节点上面都有一个颜色 $a_i$ 。颜色范围在 $[0,c)$ 中。

　　现在从树上任意一个点出发，走到任意一个点停止，走过的最短路径上的颜色依次排成一个序列。

　　问所有路径生成的序列一共有多少种。

　　$n\leq 10^5,\ \ \ 1\leq c\leq 10$

题解

　　第一次写广义 SAM ，居然只写了 20 分钟？？

　　对于每一个度为 $1$ 的节点，把该点为根的树当作一个 Trie ，并建到广义后缀自动机里面。

　　最后统计一下就可以了。

　　时间复杂度 $O(20n)$ ，空间约为 $2\times 26\times 20\times n\times (\rm {sizeof\ int})$ 。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N=100005,S=N*20*2;
struct Gragh{
	static const int M=N*2;
	int cnt,y[M],nxt[M],fst[N];
	void clear(){
		cnt=0;
		memset(fst,0,sizeof fst);
	}
	void add(int a,int b){
		y[++cnt]=b,nxt[cnt]=fst[a],fst[a]=cnt;
	}
}g;
int n,c;
int a[N],in[N];
int root,size;
struct SAM{
	int Next[10],fa,Max;
}t[S];
void init(){
	memset(t,0,sizeof t);
	root=1,size=1;
	t[0].Max=-1;
	for (int i=0;i<10;i++)
		t[0].Next[i]=1;
}
int extend(int p,int c){
	if (t[p].Next[c]&&t[p].Max+1==t[t[p].Next[c]].Max)
		return t[p].Next[c];
	int np=++size,q,nq;
	t[np].Max=t[p].Max+1;
	for (;!t[p].Next[c];p=t[p].fa)
		t[p].Next[c]=np;
	q=t[p].Next[c];
	if (t[p].Max+1==t[q].Max)
		t[np].fa=q;
	else {
		nq=++size;
		t[nq]=t[q],t[nq].Max=t[p].Max+1;
		t[q].fa=t[np].fa=nq;
		for (;t[p].Next[c]==q;p=t[p].fa)
			t[p].Next[c]=nq;
	}
	return np;
}
void dfs(int x,int pre,int p){
	for (int i=g.fst[x];i;i=g.nxt[i])
		if (g.y[i]!=pre)
			dfs(g.y[i],x,extend(p,a[g.y[i]]));
}
int main(){
	scanf("%d%d",&n,&c);
	for (int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	g.clear();
	for (int i=1,a,b;i<n;i++){
		scanf("%d%d",&a,&b);
		g.add(a,b);
		g.add(b,a);
		in[a]++,in[b]++;
	}
	init();
	for (int i=1;i<=n;i++)
		if (in[i]==1)
			dfs(i,0,extend(root,a[i]));
	LL ans=0;
	for (int i=2;i<=size;i++)
		ans+=t[i].Max-t[t[i].fa].Max;
	printf("%lld",ans);
	return 0;
}