131. 小美的树上染色（卡码网周赛第二十一期（23年美团笔试真题））

131. 小美的树上染色（卡码网周赛第二十一期（23年美团笔试真题））

题目描述

小美拿到了一棵树，每个节点有一个权值。初始每个节点都是白色。
小美有若干次操作，每次操作可以选择两个相邻的节点，如果它们都是白色且权值的乘积是完全平方数，小美就可以把这两个节点同时染红。
小美想知道，自己最多可以染红多少个节点？

输入

第一行输入一个正整数 n（1 <= n <= 10^5），代表节点的数量。第二行输入 n 个正整数 ai（1 <= ai <= 10^9），代表每个节点的权值。接下来的 n - 1 行，每行输入两个正整数 u，v（1 <= u, v <= n），代表节点 u 和节点 v 有一条边连接。

输出

输出一个整数，表示最多可以染红的节点数量。

样例输入

3
3 3 12
1 2
2 3

样例输出

2

提示

可以染红第二个和第三个节点。
请注意，此时不能再染红第一个和第二个节点，因为第二个节点已经被染红。 因此，最多染红 2 个节点。
不能使用long long类型来定义数组a，这是坑点

解题思路

题解1（C++版本）

#include<bits/stdc++.h>
using namespace std;
typedef int LL; // 这里不能使用long long类型
const int N = 1e5 + 10;

int n, dp[N][2];
LL a[N];
vector<int> edge[N];

void dfs(int u, int f){
    for(int i = 0; i < edge[u].size(); i++){
        int v = edge[u][i];
        if(v == f) continue;
        dfs(v, u);
        dp[u][0] += max(dp[v][0], dp[v][1]);
    }
    
    for(int i = 0; i < edge[u].size(); i++){
        int v = edge[u][i];
        if(v == f) continue;
        LL mul = a[u]*a[v];
        int sqr = (int)sqrt(mul);
        if(mul != 1LL*sqr * sqr) continue;
        dp[u][1] = max(dp[u][1],dp[u][0] - max(dp[v][0], dp[v][1]) + dp[v][0] + 2);
    }
}
int main(){
    scanf("%d", &n);
    for(int i = 1; i <= n; i++){
        scanf("%lld", &a[i]);
    }
    
    for(int i = 1, u, v; i < n; i++){
        scanf("%d %d", &u, &v);
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    dfs(1, 0);
    printf("%d\n", max(dp[1][0], dp[1][1]));
    return 0;
}