HDU 6133 Army Formations(线段树合并)

#include<stdio.h>
#include<string.h>
#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
const int MAXN=(int)1e5+7;
typedef long long ll;
int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
struct edge{
	int to,nxt;
}ed[MAXN<<1];
int head[MAXN],cnt;
void addedge(int u,int v){
	ed[cnt].to=v;
	ed[cnt].nxt=head[u];
	head[u]=cnt++;
}
int root[MAXN],a[MAXN];
int ls[MAXN*18],rs[MAXN*18],tol;
int num[MAXN*18];
ll sum[MAXN*18],ans[MAXN*18];int n,vsz;vector<int>ve;
void insert(int &rt,int l,int r,int x){
    if(rt==0)rt=++tol;        
	ans[rt]=sum[rt]=ve[x-1];num[rt]=1;
    if(l==r)return ;
    int mid=l+r>>1;
    if(x<=mid)insert(ls[rt],l,mid,x);
    else insert(rs[rt],mid+1,r,x);
}
int merge(int u,int v){
    if(u==0||v==0)return u|v;
    ls[u]=merge(ls[u],ls[v]);
    rs[u]=merge(rs[u],rs[v]);
    sum[u]+=sum[v];
    num[u]+=num[v];
    if(ls[u]==0&&rs[u]==0)ans[u]=(sum[u]/num[u]+sum[u])*num[u]/2;
	else ans[u]=ans[ls[u]]+ans[rs[u]]+num[rs[u]]*sum[ls[u]];
    return u;
}

void dfs(int u,int pre){
	insert(root[u],1,vsz,a[u]);
	for(int i=head[u];i!=-1;i=ed[i].nxt){
		int v=ed[i].to;
		if(v!=pre){
			dfs(v,u);
			merge(root[u],root[v]);
		}
	}
}
int main()
{
	int t;
	t=read();
	while(t--){
		ve.clear();	
		n=read();
		memset(head,-1,sizeof(head));
		memset(root,0,sizeof(root));
		for(int i=1;i<=n;i++)a[i]=read(),ve.push_back(a[i]);
		sort(ve.begin(),ve.end());
		ve.erase(unique(ve.begin(),ve.end()),ve.end());
		for(int i=1;i<=n;i++)a[i]=lower_bound(ve.begin(),ve.end(),a[i])-ve.begin()+1;
		vsz=ve.size();
		for(int i=1;i<n;i++){
			int u,v;
			u=read();v=read();
			addedge(u,v);
			addedge(v,u);
		}
		dfs(1,0);
		for(int i=1;i<=n;i++)printf("%lld ",ans[root[i]]);
		printf("\n");
		for(int i=0;i<=tol;i++)ls[i]=rs[i]=num[i]=ans[i]=sum[i]=0;
		tol=cnt=0;
	}
	return 0;
}