HDU

HDU - 6133 - Army Formations

考虑对每个结点维护一棵只含一条链的线段树，在dfs的过程中对线段树进行合并。
线段树的每个结点维护答案ans，区间元素和sum和数字个数sz。
线段树合并时：
如果是叶子结点，就是有sz[rt]个sum[rt]/sz[rt]。是个等比数列，所以ans[rt]=sum[rt]*(sz[rt]+1)/2。
如果不是叶子结点。先合并他的左右儿子。
然后ans[rt]=ans[ls[rt]]+ans[rs[rt]]+sum[rt]+sum[ls[rt]]*sz[rs[rt]]。
时间复杂度是O(nlogn) 。
空间复杂度用了动态开点还好。。

#include<bits/stdc++.h>using namespace std;typedef long long ll;const int N=1e5+7;int a[N],ls[N*9],rs[N*9],sz[N*9],b[N],m,tot,s[N*10],top;ll ans[N*9],sum[N*9],Ans[N];vector<int> adj[N];int newnode(){    if(top) return s[--top];    return ++tot;}void delnode(int x){    s[top++]=x;}void build(int &rt,int l,int r,int p){    rt=newnode();    ans[rt]=sum[rt]=b[p];    sz[rt]=1;    ls[rt]=rs[rt]=0;    if(l==r) return ;    int mid=(l+r)>>1;    if(p<=mid) build(ls[rt],l,mid,p);    else build(rs[rt],mid+1,r,p);}int Merge(int rt1,int rt2){    if(rt1==0||rt2==0) return rt1^rt2;    ls[rt1]=Merge(ls[rt1],ls[rt2]);    rs[rt1]=Merge(rs[rt1],rs[rt2]);    sum[rt1]+=sum[rt2];    sz[rt1]+=sz[rt2];    if(ls[rt1]==0&&rs[rt1]==0) ans[rt1]=sum[rt1]*(1+sz[rt1])/2;    else ans[rt1]=ans[ls[rt1]]+ans[rs[rt1]]+sum[ls[rt1]]*sz[rs[rt1]];    delnode(rt2);    return rt1;}int dfs(int u,int p){    int rt;    build(rt,1,m,a[u]);    for(int v : adj[u])    {        if(v==p) continue;        Merge(rt,dfs(v,u));    }    Ans[u]=ans[rt];    return rt;}int main(){    int T;    scanf("%d",&T);    while(T--)    {        int n;        scanf("%d",&n);        for(int i=1;i<=n;++i) scanf("%d",&a[i]),b[i]=a[i],adj[i].clear();        sort(b+1,b+1+n);        m=unique(b+1,b+1+n)-(b+1);        for(int i=1;i<=n;++i) a[i]=lower_bound(b+1,b+m+1,a[i])-b;        for(int i=1;i<n;++i)        {            int u,v;            scanf("%d%d",&u,&v);            adj[u].push_back(v);            adj[v].push_back(u);        }        top=tot=0;        dfs(1,0);        for(int i=1;i<=n;++i) printf("%I64d ",Ans[i],i==n?'\n':' ');        puts("");    }    return 0;}