CF1923E Count Paths 题解

const int N = 1000100;
const int mod = 998244353;
vector<int> adj[N];
int a[N], buc[N], n, s;
void dfs(int u, int fa) {
    int now = buc[a[u]];
    s += now;
    for (auto &v : adj[u])
        if (v != fa) {
            buc[a[u]] = 1;
            dfs(v, u);
        }
    buc[a[u]] = now + 1;
}
signed main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int T;
    cin >> T;
    while (T--) {
        cin >> n;
        for (int i = 1; i <= n; ++i) cin >> a[i];
        for (int i = 1; i <= n; ++i) adj[i].clear(), buc[i] = 0;
        for (int i = 1; i < n; ++i) {
            int x, y;
            cin >> x >> y;
            adj[x].eb(y), adj[y].eb(x);
        }
        s = 0;
        dfs(1, 0);
        cout << s << '\n';
    }
}