[ZJOI2015][LOJ2137]诸神眷顾的幻想乡(广义SAM)

题面

https://loj.ac/problem/2137

题解

前置知识

• 广义SAM https://zhuanlan.zhihu.com/p/34838533

本题给出一棵带点权的树，保证这棵树的叶子结点不超过20，求树上所有的有向路径形成的字符串中，不同字符串的数量。

发现一个事情：如果把原来的无根树，分别以每一个叶子结点作为根形成20棵有根树，那么这20棵有根树上、所有从上到下的路径（即路径上除开头外的任何一个节点是前一个节点的儿子）形成的集合的并集（不可重）就是原树中所有的有向路径形成的集合。

怎么证明呢？首先新的20棵树中，任何一条从上到下的路径都一定是原树中的某一条路径，这个比较显然；其次，对任意原树中的路径u->v，假设将原无根树以v作为根节点形成一棵有根树，那么取此树中、u的子树中的一个叶子节点w，那么在以w作为根节点的那棵有根树上，u->v就是一条从上到下的路径。

因此，只需将20棵有根树暴力合并成一个Trie，然后建立该Trie的广义SAM，此SAM能识别的字符串总个数即为答案。对所有SAM节点u，求和\(len[u]-len[fail[u]]\)即可。

代码

#include<bits/stdc++.h>

using namespace std;

#define rg register
#define In inline
#define ll long long

const int N = 1e5;
const int SN = 2e6;

In int read(){	
	int s = 0,ww = 1;
	char ch = getchar();
	while(ch < '0' || ch > '9'){if(ch == '-')ww = -1;ch = getchar();}
	while('0' <= ch && ch <= '9'){s = 10 * s + ch - '0';ch = getchar();}
	return s * ww;
}

int loc[2*SN+5];

struct SAM{
	int cnt,nx[2*SN+5][10],fail[2*SN+5];
	ll len[2*SN+5];
	void clear(){
		fail[0] = -1;
	}
	int extend(int last,int id){
		int cur = ++cnt,p;
		len[cur] = len[last] + 1;
		for(p = last;p != -1 && !nx[p][id];p = fail[p])nx[p][id] = cur;
		if(p == -1)fail[cur] = 0;
		else{
			int q = nx[p][id];
			if(len[q] == len[p] + 1)fail[cur] = q;
			else{
				int clone = ++cnt;
				len[clone] = len[p] + 1;
				fail[clone] = fail[q];
				memcpy(nx[clone],nx[q],sizeof(nx[clone]));
				fail[q] = fail[cur] = clone;
				for(;p != -1 && nx[p][id] == q;p = fail[p])nx[p][id] = clone;
			}
		}
		return cur;
	}
	In ll count(int u){
		return len[u] - len[fail[u]];
	}
	ll query(){
		ll ans = 0;
		for(rg int i = 1;i <= cnt;i++)ans += count(i);
		return ans;
	}
}S;

int head[N+5],cnt,deg[N+5];

struct node{
	int next,des;
}e[2*N+5];

In void addedge(int a,int b){
	cnt++;
	deg[a]++,deg[b]++;
	e[cnt].des = b;
	e[cnt].next = head[a];
	head[a] = cnt;
}

int w[N+5],c;

struct Trie{
	int cnt,nx[SN+5][10];
	void dfs(int cur,int u,int fa){ 
		for(rg int i = head[u];i;i = e[i].next){
			int v = e[i].des;
			if(v == fa)continue;
			if(!nx[cur][w[v]])nx[cur][w[v]] = ++cnt;
			dfs(nx[cur][w[v]],v,u);
		}
	}
	void intree(int rt){
		if(!nx[0][w[rt]])nx[0][w[rt]] = ++cnt;
		dfs(nx[0][w[rt]],rt,0);
	}
	queue<int>q;
	void build(){
		S.clear();
		q.push(0);
		while(!q.empty()){
			int u = q.front();
			q.pop();
			for(rg int i = 0;i < c;i++)if(nx[u][i]){
				int v = nx[u][i];
				loc[v] = S.extend(loc[u],i);
				q.push(v);
			}
		}
	}
}T;

int n;
vector<int>leaf;

int main(){
	n = read(),c = read();
	for(rg int i = 1;i <= n;i++)w[i] = read();
	for(rg int i = 1;i < n;i++){
		int u = read(),v = read();
		addedge(u,v);
		addedge(v,u);
	}
	for(rg int i = 1;i <= n;i++)if(deg[i] == 2)leaf.push_back(i);
	for(rg int i = 0;i < leaf.size();i++)T.intree(leaf[i]);
	T.build();
	cout << S.query() << endl;
	return 0;
}