BZOJ3926 (后缀自动机)

BZOJ3926 诸神眷顾的幻想乡

Problem :
给一个n个节点的树（n<=10^5), 每个点有一种颜色(c<=10), 询问所有点对之间路径组成字符串的种类。保证叶子节点小于等于20.
Solution :
分别从每个叶子节点开始遍历整棵树，将遍历到的字符串加入后缀自动机。
只需要修改后缀自动机的last，就可以实现添加多串。

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

const int N = 1e5 * 15 * 2;
vector <int> eg[N];
int cl[N];
int deg[N];
int n, c;

struct Suffix_Automanon
{
	int nt[N][15], fail[N], a[N];
	int p, q, np, nq;
	int tot, root;
	int newnode(int len)
	{
		for (int i = 0; i < 15; ++i) nt[tot][i] = -1;
		fail[tot] = -1; a[tot] = len;
		return tot++;
	}
	void clear()
	{
		tot = 0;
		root = newnode(0);
	}
	void insert(int ch, int &last)
	{
		p = last; last = np = newnode(a[p] + 1); 
		for (; ~p && nt[p][ch] == -1; p = fail[p]) nt[p][ch] = np;
		if (p == -1) fail[np] = root;
		else
		{
			q = nt[p][ch];
			if (a[q] == a[p] + 1) fail[np] = q;
			else
			{
				nq = newnode(a[p] + 1);
				for (int i = 0; i < 15; ++i) nt[nq][i] = nt[q][i];
				fail[nq] = fail[q]; fail[q] = fail[np] = nq;
				for (; ~q && nt[p][ch] == q; p = fail[p]) nt[p][ch] = nq;
			}
		}
	}
	void solve()
	{
		long long ans = 0;
		for (int i = 1; i < tot; ++i)
			ans += a[i] - a[fail[i]];
		cout << ans << "\n";
	}
}sam;
void dfs(int u, int fa, int now)
{
	sam.insert(cl[u], now);
	for (int i = 0; i < (int)eg[u].size(); ++i)
	{
		int v = eg[u][i];
		if (v != fa) dfs(v, u, now);
	}
}
int main()
{
	cin.sync_with_stdio(0);
	cin >> n >> c;	
	for (int i = 1; i <= n; ++i) cin >> cl[i];
	for (int i = 1; i <= n; ++i) eg[i].clear(), deg[i] = 0;
	for (int i = 1; i <  n; ++i)
	{
		int u, v; 
		cin >> u >> v;
		eg[u].push_back(v);
		eg[v].push_back(u);
		deg[u]++; deg[v]++;
	}
	sam.clear();
	for (int i = 1; i <= n; ++i)
		if (deg[i] == 1)
			dfs(i, 0, 0);
	sam.solve();
}