HDU 6133 Army Formations （树状数组, 2017 Multi-Univ Training Contest 8）

Problem

Problem Link

Though being cruel and merciless in the battlefields, the total obedience to the command hierarchy makes message delivering between Stormtroopers quite inefficient, which finally caused the escape of Luke Skywalker, the destroy of the Death Star, and the collapse of the Galactic Empire.

In particular, the hierarchy of Stormtroopers is defined by a binary tree. Everyone in the tree has at most two direct subordinates and exactly one direct leader, except the first (numbered 1) Stormtrooper, who only obey the order of the Emperor.

It has been a time-consuming task for the Stormtroopers to input messages into his own log system. Suppose that the i-th Stormtrooper has a message of length ai, which means that it costs ai time to input this message into a log system. Everyone in the hierarchy has the mission of collecting all messages from his subordinates (a direct or indirect children in the tree) and input thses messages and his own message into his own log system.

Everyone in the hierarchy wants to optimize the process of inputting. First of all, everyone collects the messages from all his subordinates. For a Stormtrooper, starting from time 0, choose a message and input it into the log system. This process proceeds until all messages from his subordinates and his own message have been input into the log system. If a message is input at time t, it will induce t units of penalty.

For every Stormtrooper in the tree, you should find the minimum penalty.

Idea

数据结构, 可以动态往一个多重集里面添加数字, 删除数字, 并查询多重集中元素的sumofsum.

维护 sum of sum 可以使用两个树状数组，一个维护当前所有点的个数（每个点有其对应大小），一个维护当前所有点值的和。

由于该树结构表明任意点最多二叉，故根据边的连通情况，处理以任意点为根的子树的 size 。建立不完全的二叉树，同时任意点的左子树大小一定小于右子树大小。

数据操作变成优先处理每个点的左子树，并利用数据结构维护当前总和 (current_ans) 。在完成左子树的处理后，将左子树的所有点值从数据结构中删除，再处理右子树，之后为处理当前子树（左子树+右子树+当前子树根节点），重新将左子树的所有点值插入数据结构中维护。

重复添加、删除左子树的节点，每个节点的添删次数不会超过 $\mathcal{O}(\log N)$ 。

简单说明：任取一个点 p ，考虑极限情况，即 p 点在所有包含 p 点的子树内都处在左子树。则包含 p 点的最小左子树大小为 1 （即仅含 p 点）。同时由于任意点的左子树大小一定小于等于右子树大小。则含 p 的此小左子树大小为 2，依次类推，1, 2, 4, 8, 16, …, $2^i (\le n)$ ，故其最多在 $\log N$ 棵左子树内，重复添删次数上限为 $\mathcal{O}(\log{N})$ 。

Code

#include<bits/stdc++.h>
using namespace std;
const int N = 100000 + 10;
struct Node {
    int a, idx, ord;
} p[N];
bool cmp_a(Node a, Node b) {    return a.a < b.a;   }
bool cmp_idx(Node a, Node b)    {   return a.idx < b.idx;   }
vector<int> g[N];
int T, n, siz[N], lft[N], rgt[N],  mxId;
void getSize(int rt, int fa) 
{
    siz[rt] = 1;
    lft[rt] = rgt[rt] = 0;
    for(int i=0;i<g[rt].size();i++)
    {
        if(g[rt][i] == fa)  continue;
        if(rgt[rt] == 0)    rgt[rt] = g[rt][i];
        else    lft[rt] = g[rt][i];
        getSize(g[rt][i], rt);
        siz[rt] += siz[ g[rt][i] ];
    }
    if(lft[rt] && siz[ lft[rt] ] > siz[ rgt[rt] ])
        swap(lft[rt], rgt[rt]);
}
long long num[N], sum[N], ans[N], cur_ans;
inline int lowbit(int x)    {   return x & -x;  }
void add(long long arr[], int x, int w) {
    for(int i=x;i<=mxId;i+=lowbit(i))
        arr[i] += w;
}
long long get(long long arr[], int x) {
    long long ret = 0;
    for(int i=x;i;i-=lowbit(i))
        ret += arr[i];
    return ret;
}
void insert(int idx, int ord) 
{
    cur_ans += (get(num, mxId) - get(num, ord)) * p[idx].a + p[idx].a;
    cur_ans += get(sum, ord);
    add(num, ord, 1);
    add(sum, ord, p[idx].a);
}
void erase(int idx, int ord)
{
    add(sum, ord, -p[idx].a);
    add(num, ord, -1);
    cur_ans -= get(sum, ord);
    cur_ans = cur_ans - (get(num, mxId) - get(num, ord)) * p[idx].a - p[idx].a;
}
void clear(int rt) 
{
    assert(rt!=0);
    erase(rt, p[rt].ord);
    if(rgt[rt]) clear(rgt[rt]);
    if(lft[rt]) clear(lft[rt]);
}
void reinsert(int rt)
{
    assert(rt!=0);
    insert(rt, p[rt].ord);
    if(lft[rt]) reinsert(lft[rt]);
    if(rgt[rt]) reinsert(rgt[rt]);
}
void calc(int rt)
{
    if(rgt[rt] == 0)    // leaf node
    {
        ans[rt] = p[rt].a;
        insert(rt, p[rt].ord);
        return;
    }

    if(lft[rt])     // prior solve left tree
    {
        calc(lft[rt]);
        clear(lft[rt]);
    }

    calc(rgt[rt]);
    if(lft[rt]) reinsert(lft[rt]);
    insert(rt, p[rt].ord);
    ans[rt] = cur_ans;
}
int main()
{
    //freopen("in", "r", stdin);
    scanf("%d", &T);
    while(T-- && scanf("%d", &n)!=EOF)
    {
        memset(num, 0, sizeof(num));
        memset(sum, 0, sizeof(sum));
        mxId = 0;
        cur_ans = 0;
        for(int i=1;i<=n;i++)
            g[i].clear();
        for(int i=1;i<=n;i++)
        {
            scanf("%d", &p[i].a);
            p[i].idx = i;
        }
        sort(p+1, p+n+1, cmp_a);
        for(int i=1;i<=n;i++)
            if(i!=1 && p[i].a == p[i-1].a)  p[i].ord = p[i-1].ord;
            else    p[i].ord = ++mxId;
        sort(p+1, p+n+1, cmp_idx);

        for(int i=1, u, v;i<n;i++)
        {   
            scanf("%d %d", &u, &v);
            g[u].push_back(v);
            g[v].push_back(u);
        }
        getSize(1, 0);
        calc(1);
        for(int i=1;i<=n;i++)
            printf("%lld ", ans[i]);
        printf("\n");
    }
}