hdu 6133---Army Formations(启发式合并+树状数组)

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
typedef long long LL;
const int N=1e5+5;
int w[N];
int val[N],id[N];
int sz[N],lson[N],rson[N],tot;
vector<int>e[N];
LL num[N],sum[N];
LL ans[N],tmp;

int calSize(int now,int pre)
{
    sz[now]=1;
    for(int j=0;j<e[now].size();j++)
    {
        int to=e[now][j];
        if(to==pre) continue;
        if(!lson[now]) lson[now]=to;
        else rson[now]=to;
        sz[now]+=calSize(to,now);
    }
    if(lson[now]&&rson[now]&&sz[lson[now]]>sz[rson[now]])
    swap(lson[now],rson[now]);
    return sz[now];
}
int lowbit(int x)
{
    return x&(-x);
}
void update(LL *a,int x,int y)
{
    while(x<=tot)
    {
        a[x]+=(LL)y;
        x+=lowbit(x);
    }
}
LL query(LL *a,int x)
{
    LL r=0;
    while(x)
    {
       r+=a[x];
       x-=lowbit(x);
    }
    return r;
}

void add(int x)
{
    tmp+=(query(num,tot)-query(num,id[x]))*(LL)w[x]+(LL)w[x];
    tmp+=query(sum,id[x]);
    update(num,id[x],1);
    update(sum,id[x],w[x]);
}
void del(int x)
{
    update(num,id[x],-1);
    update(sum,id[x],-w[x]);
    tmp-=(query(num,tot)-query(num,id[x]))*(LL)w[x]+(LL)w[x];
    tmp-=query(sum,id[x]);
}
void cle(int x)
{
    del(x);
    if(lson[x]) cle(lson[x]);
    if(rson[x]) cle(rson[x]);
}
void validate(int x)
{
    add(x);
    if(lson[x]) validate(lson[x]);
    if(rson[x]) validate(rson[x]);
}
void dfs(int now)
{
   if(!lson[now])
   {
       ans[now]=(LL)w[now];
       add(now);
       return ;
   }
   dfs(lson[now]);
   if(rson[now])
   {
       cle(lson[now]);
       dfs(rson[now]);
       validate(lson[now]);
   }
   add(now);
   ans[now]=tmp;
}
int main()
{
    int T; cin>>T;
    while(T--)
    {
        int n; scanf("%d",&n);
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&w[i]);
            val[i]=w[i];
            e[i].clear();
            num[i]=sum[i]=0;
            lson[i]=rson[i]=0;
        }
        sort(val+1,val+n+1);
        tot=unique(val+1,val+n+1)-val-1;
        for(int i=1;i<=n;i++)
        {
            id[i]=lower_bound(val+1,val+tot+1,w[i])-val;
        }
        for(int i=1;i<n;i++)
        {
            int u,v; scanf("%d%d",&u,&v);
            e[u].push_back(v);
            e[v].push_back(u);
        }
        calSize(1,-1);
        tmp=0;
        dfs(1);
        for(int i=1;i<=n;i++)
            printf("%lld ",ans[i]);
        puts("");
    }
    return 0;
}