整体思路

先定义一个概念【叶子】节点:该节点为白色节点,有且只有一个相邻节点可以一起涂成红色。若某个白色节点A存在多个相邻节点,但其中只有一个相邻的白色节点,满足双方权值的乘积是完全平方数,那么节点A也是叶子节点。

本题有一个关键的说明,也就是第一句话:小美拿到了一棵树。既然是树,那么就不存在环。因此根据贪心的思想,我们每次只取【叶子】节点和【叶子】节点相邻的节点涂成红色,然后再把红色部分剪枝。这样若干次操作之后,非【叶子】节点要么被转化成新的【叶子】节点,要么直接和【叶子】节点凑对被剪枝了。最后当【叶子】节点全部消失,那么全部涂色也就完成了,我们得到答案。

这里我使用了c++的set容器。该容器底层是红黑树,插入、查找、删除的时间复杂度都十分优秀。

具体步骤

step1:创建一个哈希表,其key是节点编号,value是一个set容器,存储该节点所有可以一起涂成红色的相邻节点。当数据输入完成,所有节点对应的set容器也就维护好了。若某个set容器的size为1,那么对应的节点key就是一个【叶子】。

step2:创建一个新的set容器leaf,用于存储所有的【叶子】节点。

step3:每次从leaf容器中取出一个【叶子】u,删除。得到u对应可以一起涂红的节点v。若v也是【叶子】节点,那么将v也从leaf容器删除,否则遍历v对应的set容器,进行剪枝。其中可能出现新的【叶子】节点,加入到leaf中。

step4:重复step3,直至容器leaf为空,输出答案。

#include <cmath>
#include <iostream>
#include <set>
#include <unordered_map>
#include <vector>
using namespace std;

bool isSquare(int u, int v) {
    long long pro = (long long)u * v;
    int n = sqrt(pro);
    return n * n == pro;
}

int main() {
    int n, u, v;
    cin >> n;
    vector<int> a(n + 1);
    unordered_map<int, set<int>> mp;
    for (int i = 1; i <= n; ++i) {
        cin >> a[i];
    }
    for (int i = 1; i < n; ++i) {
        cin >> u >> v;
        if (isSquare(a[u], a[v])) {
            mp[u].insert(v);
            mp[v].insert(u);
        }
    }
    set<int> leaf;
    for (auto &[key, value] : mp) {
        if (value.size() == 1) leaf.insert(key);
    }
    int ans = 0;
    while (!leaf.empty()) {
        u = *leaf.begin();
        leaf.erase(u);
        v = *mp[u].begin();
        if (mp[v].size() == 1) { //u和v 一对一
            leaf.erase(v);
        } else { //v包含多个
            for (const int& i : mp[v]) {
                if (i == u) continue;
                mp[i].erase(v);
                if (mp[i].size() == 0) {
                    leaf.erase(i);
                } else if (mp[i].size() == 1) {
                    leaf.insert(i);
                }
            }
        }
        ans += 2;
    }
    cout << ans;
}