「ZJOI2015」诸神眷顾的幻想乡

知识点: 广义 SAM，暴力

原题面 Loj Luogu

---

扯

萌妹子幽香有多萌

---

题意简述

给定一棵 \(n\) 个节点的树，每个节点都有一个字符 \(c_i\)。
可任意选择两个点 \(u,v\)，路径 \(u\rightarrow v\) 上的字符构成一个字符串（\(v\rightarrow u\) 构成的字符串可能不同）。
求所有可以构成的 不同的字符串的数量。
\(1\le n\le 10^5, 1\le c_i\le 10\)，最多存在 \(20\) 个叶节点。

---

分析题意

发现 任意 选择两个节点构成的所有字符串，等价于选择两个 叶节点 构成的字符串的 所有子串。

如果给定的是一条链，就非常好做了。
直接将整个串正反插入广义 SAM 中，求 \(\sum\operatorname{len}(i) - \operatorname{len}(\operatorname{link}(i))\) 即可。
但是给不得，只能来猜结论了。

考虑怎么把弯曲的路径掰♂直。
发现 任何一条路径，都会在叶节点到其他叶节点 构成的路径中出现。

题中给了一个很特别的性质：最多存在 \(20\) 个叶节点。
考虑暴力枚举选择两个 叶节点 构成的字符串，并将它们插入广义 SAM 中，求 \(\sum\operatorname{len}(i) - \operatorname{len}(\operatorname{link}(i))\) 即为答案。

复杂度大概是 \(O(\text{leaves_num}\times n)\)，可过。

---

代码实现

//知识点：广义 SAM，暴力
/*
By:Luckyblock
*/
#include <algorithm>
#include <cctype>
#include <cstdio>
#include <cstring>
#define ll long long
const int kMaxn = 1e5 + 10;
const int kMaxm = 10;
//=============================================================
ll ans;
char S[kMaxn];
bool is_leaves[kMaxn];
int leaves_num, leaves[50];
int edge_num, c[kMaxn], into[kMaxn], head[kMaxn], v[kMaxn << 1], ne[kMaxn << 1];
int node_num = 1, ch[kMaxn << 5][kMaxm], len[kMaxn << 5], link[kMaxn << 5];
//=============================================================
inline int read() {
  int f = 1, w = 0;
  char ch = getchar();
  for (; !isdigit(ch); ch = getchar())
    if (ch == '-') f = -1;
  for (; isdigit(ch); ch = getchar()) w = (w << 3) + (w << 1) + (ch ^ '0');
  return f * w;
}
void AddEdge(int u_, int v_) {
  v[++ edge_num] = v_, ne[edge_num] = head[u_], head[u_] = edge_num;
}
int Insert(int c_, int last_) {
  if (ch[last_][c_]) {
    int p = last_, q = ch[p][c_];
    if (len[p] + 1 == len[q]) return q;
    int newq = ++ node_num;
    memcpy(ch[newq], ch[q], sizeof(ch[q])); 
    len[newq] = len[p] + 1; 
    link[newq] = link[q];
    link[q] = newq; 
    for (; p && ch[p][c_] == q; p = link[p]) ch[p][c_] = newq;
    return newq;
  }
  int p = last_, now = ++ node_num;
  len[now] = len[p] + 1;
  for (; p && ! ch[p][c_]; p = link[p]) ch[p][c_] = now;
  if (! p) {link[now] = 1; return now;} 
  int q = ch[p][c_];
  if (len[q] == len[p] + 1) {link[now] = q; return now;}
  int newq = ++ node_num;
  memcpy(ch[newq], ch[q], sizeof(ch[q])); 
  link[newq] = link[q], len[newq] = len[p] + 1; 
  link[q] = link[now] = newq;
  for (; p && ch[p][c_] == q; p = link[p]) ch[p][c_] = newq;
  return now;
}
void Dfs(int u_, int fa_, int dep_) {
  S[dep_] = c[u_];
  if (is_leaves[u_] && dep_ != 1) {
    int last = 1;
    for (int i = 1; i <= dep_; ++ i) last = Insert(S[i], last);
    return ;
  }
  for (int i = head[u_]; i; i = ne[i]) {
    if (v[i] != fa_) Dfs(v[i], u_, dep_ + 1);
  }
}
//=============================================================
int main() {
  int n = read(), Marisa = read();
  for (int i = 1; i <= n; ++ i) c[i] = read();
  for (int i = 1; i < n; ++ i) {
    int u = read(), v = read(); 
    AddEdge(u, v), AddEdge(v, u);
    into[u] ++, into[v] ++;
  }
  for (int i = 1; i <= n; ++ i) {
    if (into[i] == 1) {
      leaves[++ leaves_num] = i;
      is_leaves[i] = true; 
    }
  }
  for (int i = 1; i <= leaves_num; ++ i) {
    Dfs(leaves[i], 0, 1); 
  }
  for (int i = 2; i <= node_num; ++ i) {
    ans += len[i] - len[link[i]]; 
  }
  printf("%lld\n", ans);
  return 0; 
}