1004 - Tree
题意、思路待补
#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 9;
const long long X[3] = {998244353, 1000000007, -1000000007 - 998244353};
struct lol {int x, y;} e[N << 1];
int t, n, a[N], ans, top[N], siz[N], num, vis[N], rt, h[N];
long long sum;
vector <long long> vct;
unordered_map <long long, int> g;
void ein(int x, int y) {e[++ ans].x = top[x];e[ans].y = y;top[x] = ans;
}
void dfs(int x, int fa) {siz[x] = 1; h[x] = 0;for (int i = top[x]; i; i = e[i].x) {int y = e[i].y;if (y == fa || vis[y]) continue;dfs(y, x);siz[x] += siz[y];h[x] = max(h[x], siz[y]);}h[x] = max(h[x], num - siz[x]);if (h[x] < h[rt]) rt = x;
}
void dfs2(int x, int fa, long long d, int p) {d += X[a[x]];vct.push_back(d);long long dis = d + X[p];if (dis == 0) ++ sum;if (g.find(-dis) != g.end()) sum += g[-dis];for (int i = top[x]; i; i = e[i].x) {int y = e[i].y;if (y == fa || vis[y]) continue;dfs2(y, x, d, p);}
}
void dfs1(int x) {vis[x] = 1; g.clear(); for (int i = top[x]; i; i = e[i].x) {int y = e[i].y;if (vis[y]) continue;dfs2(y, x, 0, a[x]);for (auto v : vct) ++ g[v];vct.clear();}for (int i = top[x]; i; i = e[i].x) {int y = e[i].y;if (vis[y]) continue;rt = 0; num = siz[y];dfs(y, x);dfs1(rt);}
}
int main() {scanf("%d", &n); num = n; h[0] = 1e9;for (int i = 1; i <= n; ++ i) {char c = getchar();while (c == ' ' || c == '\n') c = getchar();a[i] = c - 'a';}for (int i = 1, u, v; i < n; ++ i)scanf("%d%d", &u, &v), ein(u, v), ein(v, u);dfs(1, 0);dfs1(rt);printf("%lld", sum);return 0;
}