HackerEarth, The Grass Type （dsu on tree）

題目連結： https://www.hackerearth.com/practice/algorithms/graphs/depth-first-search/practice-problems/algorithm/the-grass-type/

題意：n個點，根為1。每個節點i都有一個值

。求使得

的無序對(u,v)個數。

思路：将問題轉化為子樹中的問題，以u為LCA的無序對(v1,v2)一定在u的子樹中，其中v1，v2要屬于u的不同兒子。是以，我們維護子樹中每個數的數量（我用的unordered_map）。在計算u子樹的答案時，統計u的輕兒子v的貢獻時，要先統計答案，再将v的子樹裡的數加上。這樣是為了避免将v的子樹的無序對（v1，v2）統計為答案，雖然

，但顯然(v1,v2)的LCA不是u。

代碼：

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 1e5+10;
int n,m,a[N];
struct edge{
	int to,nex;
}g[N<<1];
int head[N],cnt;
void add(int u,int v){
	g[cnt]=edge{v,head[u]};
	head[u]=cnt++;
}
int sz[N],son[N];
void getsz(int u,int fa){
	sz[u]=1;
	int mx=-1;
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		if(v==fa) continue;
		getsz(v,u);
		if(mx<sz[v]){
			mx=sz[v];
			son[u]=v;
		}
		sz[u]+=sz[v];
	}
}
int SON;
unordered_map<int,int> mp;
int temp[N],cntt;
ll ans=0;
void dfs1(int u,int fa,int x){
	temp[++cntt]=a[u];
	if(x%a[u]==0) ans+=mp[x/a[u]];
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		if(v==SON||v==fa) continue;
		dfs1(v,u,x);
	}	
}
void dfs(int u,int fa,bool keep){
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		if(v==fa||v==son[u]) continue;
		dfs(v,u,0);
	}
	if(son[u]!=-1){
		dfs(son[u],u,1);
		SON=son[u];
	}
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		if(v==fa||v==son[u]) continue;
		cntt=0;
		dfs1(v,u,a[u]);
		for(int i=1;i<=cntt;i++)
			mp[temp[i]]++;		
	}	
	ans+=mp[1];//注意要先統計答案，再加上a[u],否者當a[u]=1時,會多統計一個答案。
	mp[a[u]]++;
	SON=0;
	if(keep==0)
		mp.clear();
}
int main(void){
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		son[i]=head[i]=-1;
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d%d",&u,&v);
		add(u,v);
		add(v,u);
	}
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	getsz(1,0);
	dfs(1,0,0);
	cout<<ans;
	return 0;	
}