CF1923E 一个无需 DSU On Tree 的解法

在地铁上口胡了一下。不知道对不对。

考虑记录每一个点 \(i\) 离他最远的一个祖先使得祖先到 \(i\) 的路径上没有 \(a_i\)。设他为 \(\text{lst}_i\)。然后如果两个 \(u, v\) 的 \(\text{lst}\) 相等，那么这条路径就是好的。每种颜色枚举即可。

八成假了（？），欢迎 Hack。

PS：全对了，确实能过。

/*******************************
| Author:  DE_aemmprty
| Problem: E. Count Paths
| Contest: Codeforces - Educational Codeforces Round 162 (Rated for Div. 2)
| URL:     https://codeforces.com/contest/1923/problem/E
| When:    2024-02-23 23:29:35
|
| Memory:  512 MB
| Time:    2000 ms
*******************************/
#include <bits/stdc++.h>
using namespace std;

long long read() {
    char c = getchar();
    long long x = 0, p = 1;
    while ((c < '0' || c > '9') && c != '-') c = getchar();
    if (c == '-') p = -1, c = getchar();
    while (c >= '0' && c <= '9')
        x = (x << 1) + (x << 3) + (c ^ 48), c = getchar();
    return x * p;
}

const int N = 2e5 + 7;

int n;
int c[N], lst[N], f[N];
vector <int> to[N], pos[N];
vector <int> p[N];
vector <int> col;
int cnt[N];

void dfs(int u, int fa) {
    f[u] = fa;
    if (fa) p[c[fa]].push_back(u);
    lst[u] = 0;
    if (p[c[u]].size()) lst[u] = p[c[u]].back();
    for (int v : to[u])
        if (v != fa)
            dfs(v, u);
    if (fa) p[c[fa]].pop_back();
}

void solve() {
    n = read();
    for (int i = 1; i <= n; i ++) {
        c[i] = read();
        pos[i].clear();
        to[i].clear();
        p[i].clear();
    }
    for (int i = 1; i <= n; i ++)
        pos[c[i]].push_back(i);
    for (int i = 1, u, v; i < n; i ++) {
        u = read(), v = read();
        to[u].push_back(v);
        to[v].push_back(u);
    }
    dfs(1, 0); long long ans = 0;
    for (int i = 1; i <= n; i ++) {
        for (int x : pos[i]) {
            cnt[lst[x]] ++;
            col.push_back(lst[x]);
        }
        for (int x : col) {
            ans += 1ll * cnt[x] * (cnt[x] - 1) / 2;
            cnt[x] = 0;
        }
        col.clear();
    }
    for (int i = 1; i <= n; i ++)
        ans += (f[lst[i]] > 0);
    cout << ans << '\n';
}

signed main() {
    int t = 1;
    t = read();
    while (t --) solve();
    return 0;
}