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 */