D-小美的树上染色

https://ac.nowcoder.com/acm/contest/63585#question

#include<bits/stdc++.h>
using namespace std;
const int N = 1e5+10;
int h[N],ne[2*N],e[2*N],w[N];
int n,idx,k;
bool st[N];
void add(int a,int b)
{
	e[idx]=a,ne[idx]=h[b],h[b]=idx++;
}
bool check(int x,int y)
{
	x=w[x],y=w[y];
	long long t=x*y;
	long long c=sqrt(t);
	
	return c*c==t;
}
void dfs(int u,int father)
{
    //不选，跑到树的底层
	for(int i=h[u];i!=-1;i=ne[i])
	{
		int j=e[i];
		if(j==father) continue;
		
		dfs(j,u);
	}
    //从底层开始选
	for(int i=h[u];i!=-1;i=ne[i])
	{
		int j=e[i];
		if(j==father) continue;
		
		if(!st[j]&&check(j,u))
		{
			k+=2;
			st[j]=true;
			st[u]=true;
			break;
		}
	}
}
int main()
{
	cin>>n;
	memset(h,-1,sizeof h);
	for(int i=1;i<=n;i++) cin>>w[i];
	
	for(int i=1;i<n;i++)
	{
		int a,b; cin>>a>>b;
		add(a,b),add(b,a);
	}
	
	dfs(1,-1);
	cout<<k<<endl;
	return 0;
}