贪心，DFS：小美的树上染色

题目描述

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

输入描述:第一行输入一个正整数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]);
}

1. const int N=1e6+5;：定义一个常量N，表示图中顶点的最大数量。
   
2. inline int read()：定义一个内联函数read，用于读取一个整数,内联函数使运行效率高。
   
3. int a[N],f[N][2];：定义两个数组a和f，a用于存储每个顶点的值，f是一个二维数组，用于存储每个顶点的两个不同状态的值。
   
4. vector<int>g[N];：定义一个邻接表g，表示图的邻接关系。
   
5. void dfs(int x,int y)：定义一个深度优先搜索函数dfs，x是当前顶点，y是上一个顶点。
   
6. 在dfs函数中，首先初始化一个空的vector ve，两个整数k和s。
   
7. 使用循环遍历当前顶点x的所有邻接顶点，如果邻接顶点是上一个顶点y，则跳过。
   
8. 对每个邻接顶点i，递归调用dfs(i,x)。
   
9. 计算当前顶点x和邻接顶点i的乘积的平方根，如果平方根的平方等于原始乘积，则将i加入到ve中。
   
10. 更新s为当前顶点所有子树中最大值的最大值。
    
11. 更新f[x][0]为s。
    
12. 循环遍历ve中的每个顶点i，更新f[x][1]为当前f[x][1]和s减去i的较大值加上i的较小值再加2的最大值。
    
13. 读取顶点的数量n。
    
14. 循环读取每个顶点的值，并存储在数组a中。
    
15. 循环读取每条边的信息，构建图的邻接表。
    
16. 从顶点1开始调用dfs函数。
    
17. 输出f[1][0]和f[1][1]中的较大值。
    

找到树中最长的简单路径，f[x][0]表示以x为根的子树中，不包含直径的最大边权和，而f[x][1]表示包含直径的最大边权和。最终输出的是整个树中最长路径的边权和。