Luogu P3576 [POI2014]MRO-Ant colony

#include<bits/stdc++.h>
#define ll long long
using namespace std;

const int N=1e6+5;
int cnt,to[N<<1],nxt[N<<1],he[N],R[N],L[N],son[N],a[N],n,m,k;
ll ans;

inline void add(int u,int v) {
	to[++cnt]=v,nxt[cnt]=he[u],he[u]=cnt;
}

inline void dp(int u,int siz,int l,int r) {
	if((ll)siz*l>a[m]) return;
	L[u]=siz*l;
	if((ll)siz*(r+1)-1>a[m]) R[u]=a[m];
		else R[u]=siz*(r+1)-1;
}

void dfs(int fa,int u) {
	for(int e=he[u];e;e=nxt[e]) {
		int v=to[e];
		if(v!=fa) {
			dp(v,(son[v]==0?1:son[v]),L[u],R[u]);
			if(L[v]!=0||R[v]!=0) {
				dfs(u,v);
			} 
		}
	}
}


inline int left(int x) {
	int l=1,r=m,ans;
	while(l<=r) {
		int mid=l+r>>1;
		if(a[mid]>=x) ans=mid,r=mid-1;
			else l=mid+1;  
	}
	return ans;
}

inline int right(int x) {
	int l=1,r=m,ans=0;
	while(l<=r) {
		int mid=l+r>>1;
		if(a[mid]<=x) ans=mid,l=mid+1;
			else r=mid-1;
	}
	return ans;
}

int main() {
	scanf("%d%d%d",&n,&m,&k);
	for(int i=1;i<=m;i++) {
		scanf("%d",&a[i]);
	}
	sort(a+1,a+m+1);
	int rt=0,Fa=0;
	for(int i=1;i<n;i++) {
		int u,v; scanf("%d%d",&u,&v);
		add(u,v),add(v,u);
		son[u]++,son[v]++;
		if(i==1) {
			rt=u,Fa=v;
		}
	}
	for(int i=1;i<=n;i++) son[i]--;
	dp(rt,(son[rt]==0?1:son[rt]),k,k);
	dp(Fa,(son[Fa]==0?1:son[Fa]),k,k);
	dfs(Fa,rt),dfs(rt,Fa);
	for(int i=1;i<=n;i++) {
	//	if(!son[i]) printf("%d %d\n",L[i],R[i]);
		if(!son[i]&&(L[i]||R[i])) {
			ans+=right(R[i])-left(L[i])+1;
		}
	}
	printf("%lld\n",ans*k);
	return 0;
}