算法
(树上dfs) $\mathcal{O}(n)$
记 cnt[c]
表示假设当前点的颜色为 $c$,当不穿过颜色 $c$ 的顶点时,可以到达当前点的颜色 $c$ 的顶点个数
然后跑一遍 $\operatorname{dfs}$ 即可
C++ 代码
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
using ll = long long;
void solve() {
int n;
cin >> n;
vector<int> c(n);
rep(i, n) cin >> c[i];
rep(i, n) c[i]--;
vector<vector<int>> to(n);
rep(i, n-1) {
int a, b;
cin >> a >> b;
--a; --b;
to[a].push_back(b);
to[b].push_back(a);
}
ll ans = 0;
vector<int> cnt(n);
auto dfs = [&](auto& f, int v, int p=-1) -> void {
int pre = cnt[c[v]];
ans += pre;
for (int u : to[v]) {
if (u == p) continue;
cnt[c[v]] = 1;
f(f, u, v);
}
cnt[c[v]] = pre+1;
};
dfs(dfs, 0);
cout << ans << '\n';
}
int main() {
cin.tie(nullptr) -> sync_with_stdio(false);
int t;
cin >> t;
while (t--) solve();
return 0;
}