题面
题解
前置知识
本题给出一棵带点权的树，保证这棵树的叶子结点不超过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;
}