HDU - 6133 - Army Formations（线段树合并）

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;
}