2023.08.12-美团-第五题-树上染色

给定一棵树，每个节点都有一个权值以及最开始是白色。
定义操作A：
选择两个有边直接相连的节点，可以将两个节点同时染红.当且仅当他们都是白色
但是这样的题目太过简单，所以我们定义一个更复杂的操作B：
在满足操作A的条件下 两个节点的权值的乘积也需要是x∗x的形式 ，现在允许执行操作若干次操作B。问这棵树最多能够得到红色节点?

1. 动态规划树

#include"bits/stdc++.h"

using namespace std;


bool judge(long long x,long long y){
    long long i = sqrt(x*y);
    return x*y == i*i;
}


int main() {
    int n; cin>>n;
    vector<int> val(n+1);
    for(int i=1;i<=n;i++)
        cin>>val[i];
    vector<int> g[n+1];

    int from; int to;
    for(int i=0;i<n-1;i++){
        cin>>from>>to;
        g[from].push_back(to);
        g[to].push_back(from);
    }

    vector<vector<int>> memo(n+1,vector<int>(2,-1));//定义状态为第i个节点和上一节点不染和染的最大节点数
    function<int(int,bool,int)> dfs = [&](int i,bool j,int fa)->int{
        if(memo[i][j]!=-1) return memo[i][j];
        int not_dye = 0;
        int dye = (fa>0&&judge(val[i],val[fa]))?2:0;
        for (int x: g[i]) {//遍历子节点
            if (x == fa) continue;//跳过父节点
            dye+=dfs(x,0,i);//后面节点都不能染色
            int cur_not = max(dfs(x,0,i),dfs(x,1,i));//该节点可以选择染色与不染色，取最大值
            for(int y: g[i]){//其它子节点不染色，进行配对
                if(y == fa||y == x) continue;//跳过父节点和同一节点
                cur_not += dfs(y,0,i);//其它节点不染色
            }
            not_dye = max(not_dye,cur_not);
        }
        memo[i][0] = not_dye;
        memo[i][1] = dye;
        return memo[i][j];   
    };

    cout<<dfs(1,0,0);
    return 0;
};