题解 | #小美的树上染色#

我们可以直接dfs这棵树，在回溯的时候判断，如果当前节点和其父亲是满足条件的且都是白色节点，那么一定可以将其染色，因为这时一定代表着当前节点 u 不存在儿子 v ，使得 u v 满足条件，不然 u 已经是红色了，这样贪心的染色一定可以染出最大答案

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 2e5 + 5;
int __t = 1, n, u, v, w[N], vis[N], ans = 0;
vector<int> g[N];
void dfs(int u, int d) {
    for (auto v : g[u]) {
        if (v == d)
            continue;
        dfs(v, u);
    }
    if (vis[u] || vis[d] || d == 0)
        return;
    int num = w[d] * w[u];
    if (sqrt(num) == (int)sqrt(num)) {
        vis[u] = vis[d] = 1;
        ans += 2;
    }
    return;
}
void solve() {
    cin >> n;
    for (int i = 1; i <= n; i++)
        cin >> w[i];
    for (int i = 1; i <= n - 1; i++) {
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1, 0);
    cout << ans << '\n';
    return;
}
int32_t main() {
#ifdef ONLINE_JUDGE
    ios::sync_with_stdio(false);
    cin.tie(0);
#endif
    // cin >> __t;
    while (__t--)
        solve();
    return 0;
}