整体思路
先定义一个概念【叶子】节点:该节点为白色节点,有且只有一个相邻节点可以一起涂成红色。若某个白色节点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; }