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

D题解 | #小美的树上染色# 经典的树形DP

1、以节点1为根节点建树（任何一个节点为根节点都行）；

2、dp[i][0] 表示以节点i为子树根节点时，节点i没有被染色时整个i节点子树可以被染色的最多节点数量；同样 dp[i][1] 表示以节点i为子树根节点时，节点i被染色时整个i节点子树可以被染色的最多节点数量；

3、dp[i][0] 为i节点的所有下级子节点j的最大值总和,为 dp[i][0]=∑max(dp[j][0],dp[j][1])（j为i的下级子节点）； dp[i][1] 为i节点与某一子节点j能够配对的最大值，dp[i][1] = max(dp[i][1], 2+dp[i][0]-max(dp[j][0],dp[j][1]))。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <map>
#include <unordered_set>
#include <vector>
#include <set>
#include <queue>
#include <iostream>
using namespace std;
vector<int> vc[100005];
long long a[100005];
int dp[100005][2];
int b[100005];
void bfs(int x){
    int i, y=0, all=0;
    vector<int> v1;
    for(i=0;i<vc[x].size();++i){
        if (b[vc[x][i]] ==0){
            v1.push_back(vc[x][i]);b[vc[x][i]] =1;
        }
    }
    for(i=0;i<v1.size();++i) bfs(v1[i]);
    vector<int> v2;
    for(i=0;i<v1.size();++i) all+=max(dp[v1[i]][0], dp[v1[i]][1]);
    dp[x][0] = all;
    long long t,k;
    for(i=0;i<v1.size();++i){
        t = a[x] * a[v1[i]]; k=sqrt(t);
        if (k*k ==t){
            y = max(y, 2+all-max(dp[v1[i]][0], dp[v1[i]][1]) + dp[v1[i]][0]);
        }
    }
    dp[x][1] = y;
    return;
}
void my_ans(){
    int i,j,n,t,x=0,y;
    cin>>n;
    for(i=1;i<=n;++i) cin>>a[i];
    for(i=1;i<n;++i){
        cin>>x>>y; vc[x].push_back(y); vc[y].push_back(x);
    }
    memset(dp,0,sizeof(dp));
    memset(b,0,sizeof(b));
    b[1] = 1;
    bfs(1);
    cout<<max(dp[1][0], dp[1][1])<<endl;
    return;
}
int main() {
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    long long t=1,i,j;
    //scanf("%d",&t);
    //cin>>t;
    while(t>0){
        --t;my_ans();
    }
    return 0;
}
// 64 位输出请用 printf("%lld")