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

D题 可以当成二分图最大匹配问题来做 这里用的是匈牙利算法 (但是没有dp快

#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 10;
int n,match[N],w[N];
vector<int> v[N];
bool vis[N];
bool check(int x)
{
    int temp = sqrt(x);
    if(temp*temp == x) return true;
    return false;
}
bool dfs(int u)
{
    vis[u] = true;
    for(auto t : v[u])
    {
        if(vis[t]) continue;
        vis[t] = true;
        if(match[t] == -1 || dfs(match[t]))
        {
            match[t] = u;
            match[u] = t;
            return true;
        }
    }
    return false;
}
int main()
{
    cin >> n;
    memset(match,-1,sizeof match);
    for(int i = 1; i <= n; i++)scanf("%d",&w[i]);
    for(int i = 1; i < n; i++)
    {
        int a,b;
        scanf("%d %d",&a,&b);
        if(!check(w[a]*w[b])) continue;
        v[a].push_back(b);
        v[b].push_back(a);
    }
    
    int res = 0;
    for(int i = 1; i <= n; i++)
    {
        if(match[i] != -1) continue;
        memset(vis,0,sizeof vis);
        if(dfs(i)) res++;
    }
    cout << res*2 << endl;
    return 0;
}