题解：B4189 [中山市赛 2024] 树上开花

#include<bits/stdc++.h>
#define rd read()
#define gc pa == pb && (pb = (pa = buf) + fread(buf, 1, 100000, stdin), pa == pb) ? EOF : *pa++
#define rd read()
#define md 998244353
#define fr first
#define se second
#define N 500005
#define int long long
#define V 500000
#define inf LONG_LONG_MAX
using namespace std;
static char buf[100000], * pa(buf), * pb(buf);
int n,ans,a[N];
vector<int>v[N];
inline int read(){
	register int x=0,ss=1,s=gc;
	while(!isdigit(s)&&s!='-')s=gc;
	if(s=='-')ss=-1,s=gc;
	while(isdigit(s))x=(x<<1)+(x<<3)+(s^48),s=gc;
	return ss*x;
}
int tot;
struct{int c,ls,rs;}tr[N*50];
inline void add(int &k,int l,int r,int x){
	if(!k) k=++tot;
	if(l==r) return tr[k].c++,void();
	int mid=l+r>>1;
	if(x<=mid) add(tr[k].ls,l,mid,x);
	else add(tr[k].rs,mid+1,r,x);
}
inline void merge(int &x,int &y,int l,int r){
	if(!y) return;
	if(!x) return x=y,void();
	if(l==r) return tr[x].c+=tr[y].c,void();
	int mid=l+r>>1;
	merge(tr[x].ls,tr[y].ls,l,mid);
	merge(tr[x].rs,tr[y].rs,mid+1,r);
}
inline int ask(int k,int l,int r,int x){
	if(!k) return 0;
	if(l==r) return tr[k].c;
	int mid=l+r>>1;
	if(x<=mid) return ask(tr[k].ls,l,mid,x);
	else return ask(tr[k].rs,mid+1,r,x);
}
inline int dfs(int x,int fa){
	int trt,rt=0;
	for(int t:v[x]){
		if(t==fa) continue;
		trt=dfs(t,x);ans+=ask(rt,1,n,a[x])*ask(trt,1,n,a[x]);
		merge(rt,trt,1,n);
	}
	ans+=ask(rt,1,n,a[x]);
	for(int i=1;i*i<=a[x];i++){
		if(a[x]%i==0){
			if(i*i!=a[x]) add(rt,1,n,a[x]/i);
			add(rt,1,n,i);
		}
	}
	return rt;
}
signed main(){
//	freopen("P4556_12.in","r",stdin);
//	freopen("sosomst.out","w",stdout);
	n=rd;
	for(int i=1;i<=n;i++) a[i]=rd;
	for(int i=1,x,y;i<n;i++){
		x=rd;y=rd;
		v[x].push_back(y);v[y].push_back(x);
	}
	dfs(1,0);cout<<ans*2+n;
	return 0;
}