牛客15天刷题ZT50_小美的树上染色

小美的树上染色

这是题面

思路

题目的限制条件有两个

1. 节点需要是相邻的

2. 都还没被访问过并且乘积是一个完全平方数

然后求最多能访问多少个节点

如果从根开始考虑的话，可能不是很好想，因为跟可能有很多儿子，和哪个儿子结合呢？好像不好想

那我们可以试着从叶子节点开始考虑，因为叶子节点能够结合的只有它们的父亲

$dfs$遍历，回溯的时候看这个儿子能不能和父亲结合，如果能结合肯定是要结合的，如果父亲能和别的儿子结合，那么结果和现在一样，如果不能，现在就是最优，所以能结合就要结合，然后用$vis$数组标记访问过

代码

神奇的代码

#include <bits/stdc++.h>
#define endl '\n'
#define int long long
const int maxn = 1e5 + 5;

std::vector<int> tree[maxn];
int vis[maxn];
int val[maxn];
int ans = 0;
void dfs(int s, int fa)
{
    for (int i = 0; i < tree[s].size(); i++)
    {
        int t = tree[s][i];
        if(t == fa) continue;
        dfs(t, s);
        if (vis[t]) continue;
        if (vis[s]) continue;
        int tmp = val[s] * val[t];
        int sq = sqrt(tmp);
        if (sq * sq == tmp)
        {
            ans += 2;
            vis[t] = 1;
            vis[s] = 1;
        }
    }
}

void solve()
{
    int n = 0;
    std::cin >> n;
    for (int i = 1; i <= n; i++)
    {
        std::cin >> val[i];
    }
    int u = 0, v = 0;
    for (int i = 1; i <= n - 1; i++)
    {
        std::cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs(1, 0);
    std::cout << ans << endl;
}

signed main()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr); std::cout.tie(nullptr);
    int t = 1;
    // std::cin >> t;
    while(t--)
    {
        solve();
    }
    return 0;
}