914E - Palindromes in a Tree - CodeForces Solution

/* Oon commento looloo bord */

#include <bits/stdc++.h>

using namespace std;

#define int long long

const int N = 200'000 + 7, Z = 22;

int n, sz[N], ans[N], cnt[1 << Z], val[N], tmp[N], cen, maj;
string s;
vector<int> adj[N], change;
bool mark[N];

void read_input() {
    cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        --u;
        --v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    cin >> s;
}

int dfs_sz(int v, int par = -1) {
    sz[v] = 1;
    for (int u: adj[v])
        if (u != par && !mark[u])
            sz[v] += dfs_sz(u, v);
    return sz[v];
}

void dfs(int v, int par = -1) {
    val[v] = 1 << (s[v] - 'a');
    val[v] ^= val[par];
    if (!cnt[val[v]])
        change.push_back(val[v]);
    cnt[val[v]]++;
    for (int u: adj[v])
        if (u != par && !mark[u])
            dfs(u, v);
}

void modify(int v, int x, int par = -1) {
    cnt[val[v]] += x;
    for (int u: adj[v])
        if (u != par && !mark[u])
            modify(u, x, v);
}

int calc(int mask) {
    int res = (__builtin_popcount(mask) <= 1);
    maj += res;
    mask ^= 1 << (s[cen] - 'a');
    res += cnt[mask];
    for (int i = 0; i < 20; i++)
        res += cnt[mask ^ (1 << i)];
    return res;
}

void add(int v, int par = -1) {
    tmp[v] = calc(val[v]);
    for (int u: adj[v])
        if (u != par && !mark[u]) {
            add(u, v);
            tmp[v] += tmp[u];
        }
    ans[v] += tmp[v];
}

int find_centroid(int v, int mx_size, int par = -1) {
    for (int u: adj[v])
        if (!mark[u] && u != par && sz[u] > mx_size)
            return find_centroid(u, mx_size, v);
    return v;
}

void clean() {
    for (int x: change)
        cnt[x] = 0;
    change.clear();
}

void decomp(int v) {
    v = find_centroid(v, dfs_sz(v) / 2);
    cen = v;
    mark[v] = true;
    val[v] = 1 << (s[v] - 'a');
    int sum = 0;
    maj = 0;
    for (int u: adj[v])
        if (!mark[u])
            dfs(u, v);
    for (int u: adj[v])
        if (!mark[u]) {
            modify(u, -1);
            add(u);
            sum += tmp[u];
            modify(u, 1);
        }
    ans[v] += (sum - maj) / 2 + maj;
    clean();
    for (int u: adj[v])
        if (!mark[u])
            decomp(u);
}

int32_t main() {
    read_input();
    decomp(0);
    for (int v = 0; v < n; v++)
        cout << ans[v] + 1 << ' ';
    cout << endl;
}