8.12美团笔试

第五题，树上染色问题。每个节点有一个值，初始状态树上节点全部为白色，每次只能选取两个相邻的节点染成红色，并且这两个节点的值的乘积是完全平方数。求最多能将多少个节点染红

#include <iostream>
#include <bits/stdc++.h>
const int maxn=1e5+100;
#define int long long 
using namespace std;
int dp[maxn][2];//dp[i][j], i,j 这个状态下能染红最多的节点
int a[maxn];
vector<int> g[maxn];
int n;
inline bool check(int x,int y){
    int mid= sqrt(x*y);
    if(mid*mid==x*y) return true;
    return false;
}
int trans_table[maxn]={0};
int find_max_trans(const vector<int> &trans){
    int val=0;
    for(int v:trans){
        if(dp[v][1]>dp[v][0]) trans_table[v]=1;
        else trans_table[v]=0;
        val+=dp[v][trans_table[v]];
    }
    int mx=0;
    for(int v:trans){
        //选择v染红
        if(trans_table[v]==1){
            val-=dp[v][1];
            val+=dp[v][0];
        }
        mx=max(mx,val+1);

        if(trans_table[v]==1){//还原
            val+=dp[v][1];
            val-=dp[v][0];
        }
    }
    return mx;

}
void dfs(int u,int fa){
    vector<int> trans;
    for(int v:g[u]){
        if(v==fa) continue;
        dfs(v,u);

        if(check(a[u],a[v])){
            trans.push_back(v);
        }else{
            dp[u][1]+=max(dp[v][0],dp[v][1]);//无限制转移
        }
        dp[u][0]+=max(dp[v][0],dp[v][1]);//无限制转移
    }
    if(trans.size()==0){//u无法染成红色
        dp[u][1]=0;
    }
    else dp[u][1]+=find_max_trans(trans)+1;

}
signed main() {
//	int mid=sqrt(999999999ll*999999999ll);
//	cout<<mid<<endl;
    cin>>n;
    memset(dp,0,sizeof(dp));
    for(int i=1;i<=n;i++) cin>>a[i];
    for(int i=1;i<=n-1;i++){
        int u,v;cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,-1);
    int ans=max(dp[1][0],dp[1][1]);
    cout<<ans<<endl;
    
}
// 64 位输出请用 printf("%lld")
/*
5
1 1 1 1 1 
1 2
1 3
1 4
1 5
*/