【洛谷P3346】诸神眷顾的幻想乡

题目

题目链接：https://www.luogu.com.cn/problem/P3346
幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的 \(2600\) 岁生日，无数幽香的粉丝到了幽香家门前的太阳花田上来为幽香庆祝生日。
粉丝们非常热情，自发组织表演了一系列节目给幽香看。幽香当然也非常高兴啦。
这时幽香发现了一件非常有趣的事情，太阳花田有 \(n\) 块空地。
在过去，幽香为了方便，在这 \(n\) 块空地之间修建了 \(n-1\) 条边将它们连通起来。
也就是说，这 \(n\) 块空地形成了一个树的结构。
有 \(n\) 个粉丝来到了太阳花田上。
为了表达对幽香生日的祝贺，他们选择了 \(c\) 种颜色的衣服，每种颜色恰好可以用一个 \(0\) 到 \(c-1\) 之间的整数来表示。
并且每个人都站在一个空地上，每个空地上也只有一个人。
这样，整个太阳花田就花花绿绿了。幽香看到了，感觉也非常开心。
粉丝们策划的一个节目是这样的，选中两个粉丝 \(A\) 和 \(B\)（\(A\) 和 \(B\) 可以相同），然后A所在的空地到B所在的空地的路径上的粉丝依次跳起来（包括端点）。
这样幽香就能看到一个长度为 \(A\) 到 \(B\) 之间路径上的所有粉丝的数目（包括 \(A\) 和 \(B\)）的颜色序列。
一开始大家打算让人一两个粉丝（注意：\(A,B\) 和 \(B,A\) 是不同的，他们形成的序列刚好相反）都来一次。
但是有人指出这样可能会出现一些一模一样的颜色序列，会导致审美疲劳。
于是他们想问，在这个树上，一共有多少可能的不同的颜色序列幽香可以看到呢？
由于太阳花田的结构比较特殊，只与一个空地相邻的空地数量不超过 \(20\) 个。
\(n\leq 10^5,c\leq 10\)。

思路

树上路径本质不同子串数量。考虑 SAM。
但是 SAM 只支持一个字符串的操作，所以我们考虑广义 SAM。
广义 SAM 需要我们把多个字符串建成 Trie，然后在遍历 Trie 依次插入字符。题目中给出的树有点类似，但是题目中路径不一定是一直从上往下的，可能会出现转折。
注意到叶子节点数量只有 \(20\)，我们发现任意一条路径，总可以在某一个叶子节点为根的时候，变成一条从上往下的路径，所以我们可以枚举叶子节点作为根，分别当做 Trie 插入一遍广义 SAM，这样所有子串都会在 SAM 中出现了。
最后枚举所有节点把它的 \(\text{len}\) 和父亲减一下求和即可。
时间复杂度 \(O(kn)\)，其中 \(k\) 是叶子的数量。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=200010,M=11;
int n,m,tot,a[N],head[N],deg[N],last[N];
ll ans;

struct edge
{
	int next,to;
}e[N];

void add(int from,int to)
{
	e[++tot]=(edge){head[from],to};
	head[from]=tot;
}

struct SAM
{
	int tot,last,fa[N*20],len[N*20],ch[N*20][M];
	SAM() { tot=last=1; }
	
	int ins(int c,int p)
	{
		if (ch[p][c])
		{
			int q=ch[p][c];
			if (len[q]==len[p]+1) last=q;
			else
			{
				int nq=++tot; last=nq;
				len[nq]=len[p]+1; fa[nq]=fa[q];
				for (int i=0;i<m;i++) ch[nq][i]=ch[q][i];
				fa[q]=nq;
				for (;p && ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
			}
		}
		else
		{
			int np=++tot; last=np;
			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[q]==len[p]+1) fa[np]=q;
				else
				{
					int nq=++tot;
					len[nq]=len[p]+1; fa[nq]=fa[q];
					for (int i=0;i<m;i++) ch[nq][i]=ch[q][i];
					fa[q]=fa[np]=nq;
					for (;p && ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
				}
			}
		}
		return last;
	}
}sam;

void dfs(int x,int fa)
{
	last[x]=sam.ins(a[x],last[fa]);
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa) dfs(v,x);
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	for (int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
		deg[x]++; deg[y]++;
	}
	last[0]=1;
	for (int i=1;i<=n;i++)
		if (deg[i]==1) dfs(i,0);
	for (int i=1;i<=sam.tot;i++)
		ans+=sam.len[i]-sam.len[sam.fa[i]];
	printf("%lld",ans);
	return 0;
}