[Luogu P3346] [BZOJ 3926] [ZJOI2015]诸神眷顾的幻想乡

洛谷传送门

BZOJ传送门

题目描述

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的$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,c$。表示空地数量和颜色数量。 第二行有$n$个$0$到$c-1$之间，由空格隔开的整数，依次表示第$i$块空地上的粉丝的衣服颜色。（这里我们按照节点标号从小到大的顺序依次给出每块空地上粉丝的衣服颜色）。 接下来$n-1$行，每行两个正整数$u,v$，表示有一条连接空地$u$和空地$v$的边。

输出格式：

一行，输出一个整数，表示答案。

输入输出样例

输入样例#1：

7 3
0 2 1 2 1 0 0
1 2
3 4
3 5
4 6
5 7
2 5

输出样例#1：

30

说明

对于所有数据，$1\le n\le 100000,1\le c\le 10$。

对于15%的数据，$n\le 2000$。

另有5%的数据，所有空地都至多与两个空地相邻。

另有5%的数据，除一块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。

另有5%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。

解题分析

我们发现这道题有个妙妙的性质：树的叶节点数量不超过$20$。 我们又可以发现树上的有向路径都是从某个叶节点开始$DFS$的得到的， 那么我们显然可以从每个叶节点出发$DFS$整棵树同时插入广义后缀自动机， 最后每个点的$len$减去其$parent$节点的$len$就是当前点代表的本质不同的子串的数量。

广义后缀自动机的节点指的是公共子串， 所以时刻注意是否满足转移条件， 添加辅助点。

注意插入时需要保存当前的$last$的位置， 向下$DFS$是从这个点继续插入的。

代码如下：

#include 
#include 
#include 
#include 
#include 
#include 
#define R register
#define IN inline
#define W while
#define gc getchar()
#define SIZ 4000500
#define MX 100050
template <class T>
IN void in(T &x)
{
	x = 0; R char c = gc;
	for (; !isdigit(c); c = gc);
	for (;  isdigit(c); c = gc)
	x = (x << 1) + (x << 3) + c - 48;
}
int to[SIZ][11], par[SIZ], len[SIZ], head[MX], deg[MX], val[MX];
int dot, last, cur, cnt, arr, typ;
long long ans;
struct Edge {int to, nex;} edge[MX << 1];
IN void add(R int from, R int to)
{edge[++arr] = {to, head[from]}, head[from] = arr;}
namespace SAM
{
	IN int insert(R int pos, R int id)
	{
		R int now = pos, tar;
		if(to[now][id])
		{
			tar = to[now][id];
			if(len[tar] == len[now] + 1) return last = tar;
			cur = ++cnt; len[cur] = len[now] + 1;
			par[cur] = par[tar], par[tar] = cur;
			std::memcpy(to[cur], to[tar], sizeof(to[cur]));
			for (; (~now) && to[now][id] == tar; now = par[now]) to[now][id] = cur;
			return last = cur;
		}
		else
		{
			cur = ++cnt; len[cur] = len[now] + 1; last = cur;
			for (; (~now) && !to[now][id]; now = par[now]) to[now][id] = cur;
			if(now < 0) return par[cur] = 0, last;
			tar = to[now][id];
			if(len[tar] == len[now] + 1) return par[cur] = tar, last;
			int nw = ++cnt; len[nw] = len[now] + 1;
			par[nw] = par[tar], par[tar] = par[cur] = nw;
			std::memcpy(to[nw], to[tar], sizeof(to[nw]));
			for (; (~now) && to[now][id] == tar; now = par[now]) to[now][id] = nw;
			return last;
		}
	}
	void DFS(R int now, R int fa, R int pos)
	{
		int rt = SAM::insert(pos, val[now]);//记住当前节点的位置
		for (R int i = head[now]; i; i = edge[i].nex)
		{
			if(edge[i].to == fa) continue;
			DFS(edge[i].to, now, rt);
		}
	}
	IN void calc()
	{
		for (R int i = 1; i <= cnt; ++i) ans += len[i] - len[par[i]];
		printf("%lld", ans);
	}
}
int main(void)
{
	int a, b; par[0] = -1;
	in(dot), in(typ);
	for (R int i = 1; i <= dot; ++i) in(val[i]);
	for (R int i = 1; i < dot; ++i)
	{
		in(a), in(b), add(a, b), add(b, a);
		deg[a]++, deg[b]++;
	}
	for (R int i = 1; i <= dot; ++i)
	if(deg[i] == 1) last = 0, SAM::DFS(i, 0, 0);
	SAM::calc();
}