题目描述
小美拿到了一棵树，每个节点有一个权值。初始每个节点都是白色。小美有若干次操作，每次操作可以选择两个相邻的节点，如果它们都是白色且权值的乘积是完全平方数，小美就可以把这两个节点同时染红。小美想知道，自己最多可以染红多少个节点？
输入描述:第一行输入一个正整数n，代表节点的数量。 第二行输入n个正整数ai，代表每个节点的权值。 接下来的n−1行，每行输入两个正整数u,v，代表节点u和节点v有一条边连接。1≤n≤10^5 1≤ai≤10^9 ,1≤u,v≤n。
输出描述:输出一个整数，表示最多可以染红的节点数量。
输入
3 3 3 12 1 2 2 3
输出
2 说明：可以染红第二个和第三个节点。请注意，此时不能再染红第一个和第二个节点，因为第二个节点已经被染红。因此，最多染红 2 个节点。
#include<bits/stdc++.h>
#include<iostream>
using namespace std;
const int N=1e6+5;
inline int read(){
int x=0;
char c=getchar();
while(c<'0'||c>'9')
c=getchar();
while(c>='0'&&c<='9')
x=(x<<3)+(x<<1)+(c^48),c=getchar();
return x;
}
int a[N],f[N][2];
vector<int>g[N];
void dfs(int x,int y){
vector<int>ve;
int k=0,s=0;
for(int i:g[x]){
if(i==y)
continue;
dfs(i,x);
int sq=sqrt(a[x]*a[i]);
if(sq*sq==a[x]*a[i])ve.push_back(i);
s+=max(f[i][1],f[i][0]);
}
f[x][0]=s;
for(int i:ve){
f[x][1]=max(f[x][1],s-max(f[i][1],f[i][0])+f[i][0]+2);
}
}
int main(){
int n=read();
for(int i=1;i<=n;++i)a[i]=read();
for(int i=1;i<n;++i){
int u=read(),v=read();
g[u].push_back(v);
g[v].push_back(u);
}
dfs(1,0);
cout<<max(f[1][0],f[1][1]);
}
-
const int N=1e6+5;：定义一个常量N，表示图中顶点的最大数量。 -
inline int read()：定义一个内联函数read，用于读取一个整数,内联函数使运行效率高。 -
int a[N],f[N][2];：定义两个数组a和f，a用于存储每个顶点的值，f是一个二维数组，用于存储每个顶点的两个不同状态的值。 -
vector<int>g[N];：定义一个邻接表g，表示图的邻接关系。 -
void dfs(int x,int y)：定义一个深度优先搜索函数dfs，x是当前顶点，y是上一个顶点。 -
在
dfs函数中，首先初始化一个空的vectorve，两个整数k和s。 -
使用循环遍历当前顶点
x的所有邻接顶点，如果邻接顶点是上一个顶点y，则跳过。 -
对每个邻接顶点
i，递归调用dfs(i,x)。 -
计算当前顶点
x和邻接顶点i的乘积的平方根，如果平方根的平方等于原始乘积，则将i加入到ve中。 -
更新
s为当前顶点所有子树中最大值的最大值。 -
更新
f[x][0]为s。 -
循环遍历
ve中的每个顶点i，更新f[x][1]为当前f[x][1]和s减去i的较大值加上i的较小值再加2的最大值。 -
读取顶点的数量
n。 -
循环读取每个顶点的值，并存储在数组
a中。 -
循环读取每条边的信息，构建图的邻接表。
-
从顶点1开始调用
dfs函数。 -
输出
f[1][0]和f[1][1]中的较大值。
找到树中最长的简单路径，f[x][0]表示以x为根的子树中，不包含直径的最大边权和，而f[x][1]表示包含直径的最大边权和。最终输出的是整个树中最长路径的边权和。