Cod sursă (job #143334)

/*
 * Vlad Ionescu - Universitatea Politehnica Bucuresti
 * Rafaela - O(M * log^2N)
 */
#include <iostream>
#include <fstream>
#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>
#include <set>
using namespace std;
#define maxn 200100

int N, Q;
vector<int> G[maxn];
multiset<int> sons[maxn];

int number_of_chains, no_of_mans;
int initial[maxn], initial_sum_subarb[maxn], updated_initial_sum_subarb[maxn], sum_mans[maxn], father[maxn];
bool extreme[maxn];
int viz[maxn], nodes_subarb[maxn], chain[maxn], pos_in_chain[maxn], nodes_in_chain[maxn], start_node[maxn], son_in_chain[maxn];

int ** sum_heavy_aint;

void dfs_max_subarb(int root, int fath) {
    viz[root] = 1;
    nodes_subarb[root] = 1;
    initial_sum_subarb[root] = initial[root];
    updated_initial_sum_subarb[root] = initial[root];
    father[root] = fath;

    for(int i=0; i<G[root].size(); i++) {
        if(!viz[ G[root][i] ]) {
            dfs_max_subarb( G[root][i], root );

            nodes_subarb[root] += nodes_subarb[ G[root][i] ];
            initial_sum_subarb[root] += initial_sum_subarb[ G[root][i] ];
            updated_initial_sum_subarb[root] += updated_initial_sum_subarb[ G[root][i] ];
        }
    }
}

void do_heavy_path(int node, bool is_extreme, int chain_no, int pos) {
    extreme[node] = is_extreme;
    chain[node] = chain_no;
    pos_in_chain[node] = pos;
    nodes_in_chain[chain_no] ++;

    if(is_extreme) {
        start_node[chain_no] = node;
    }

    int max_son = -1;
    int which_son = -1;

    // select son with maximum nodes in its subtree
    for(int i=0; i<G[node].size(); i++) {
        if(G[node][i] != father[node]) {
            if(nodes_subarb[ G[node][i] ] > max_son) {
                max_son = nodes_subarb[ G[node][i] ];
                which_son = G[node][i];
            }
        }
    }

    if(which_son != -1) {
        // continue heavy path through this son
        son_in_chain[node] = which_son;
        do_heavy_path(which_son, false, chain_no, pos + 1);

        // create new chain starting in all other sons
        for(int i=0; i<G[node].size(); i++) {
            if(G[node][i] != father[node] && G[node][i] != which_son) {
                sum_mans[node] += initial_sum_subarb[ G[node][i] ];
                sons[node].insert(initial_sum_subarb[ G[node][i] ]);

                number_of_chains ++;
                do_heavy_path(G[node][i], true, number_of_chains, 1);
            }
        }
    }
}

void create_heavy_path(int root) {
    dfs_max_subarb(root, 0);

    number_of_chains = 1;

    do_heavy_path(root, true, 1, 1);

    // initialize memory
    sum_heavy_aint = (int **)malloc(number_of_chains * sizeof(int *));
    for(int i=1; i<=number_of_chains; i++) {
        sum_heavy_aint[i] = (int *)calloc(4 * nodes_in_chain[i] + 1, sizeof(int));
    }

    // initial, all sums added are 0
    // (and it's ok because calloc())
}

int global_aint_answer;

void update_heavy_aint(int which, int node, int begin, int end, int finish, int value) {
    // [1, finish]
    // end = no of elements of this aint
    // add +value to [1, finish] interval

    if(1 <= begin && end <= finish) {
        sum_heavy_aint[which][node] += value;

        return;
    }

    int mid = ((begin + end) >> 1);

    update_heavy_aint(which, 2*node, begin, mid, finish, value);

    if(mid < finish) {
        update_heavy_aint(which, 2*node+1, mid+1, end, finish, value);
    }
}

void query_heavy_aint(int which, int node, int begin, int end, int finish) {
    // [1, finish]
    // end = no of elements of this aint
    // get sum for [1, finish] interval
    global_aint_answer += sum_heavy_aint[which][node];

    if(begin == end) {
        // global_aint_answer += sum_heavy_aint[which][node];

        return;
    }

    int mid = ((begin + end) >> 1);

    if(finish <= mid) {
        query_heavy_aint(which, 2*node, begin, mid, finish);
    }
    else {
        query_heavy_aint(which, 2*node+1, mid+1, end, finish);
    }
}

void update(int node, int value) {
    // update aint
    int finish = pos_in_chain[node];

    // add "value" to (1, finish)
    update_heavy_aint(chain[node], 1, 1, nodes_in_chain[ chain[node] ], finish, value);


    // then proceed to the next chain
    int starting_node = start_node[ chain[node] ];
    int fath = father[starting_node];

    // update set for father of starting_node
    if(fath != 0) {
        sons[fath].erase( sons[fath].find(updated_initial_sum_subarb[starting_node]) );

        sum_mans[fath] += value;
        updated_initial_sum_subarb[starting_node] += value;

        sons[fath].insert( updated_initial_sum_subarb[starting_node] );
    }

    // update next chain
    if(fath != 0) {
        update(fath, value);
    }
}

int query(int node) {
    int ret = 0;
    int sum_mans_subarb = 0;

    // first step - check how much is added to son_in_chain[node]
    if(son_in_chain[node] != 0) {
        int pos = pos_in_chain[ son_in_chain[node] ];
        int which = chain[node];

        int first_ret = initial_sum_subarb[ son_in_chain[node] ];

        global_aint_answer = 0;
        query_heavy_aint(which, 1, 1, nodes_in_chain[which], pos);
        first_ret += global_aint_answer;


        sum_mans_subarb += first_ret;

        ret = max(ret, first_ret);
    }

    // second step - find maximum from all other sons
    int how_many_sons = G[node].size();
    if(father[node] != 0) {
        how_many_sons --;
    }
    how_many_sons --; // because one son continues the chain

    if(how_many_sons > 0) {
        multiset<int>::iterator it = sons[node].end();
        it --;

        ret = max(ret, (*it));
    }

    // third step - check edge connecting with the father
    sum_mans_subarb += sum_mans[node];
    sum_mans_subarb += initial[node];

    ret = max(ret, no_of_mans - sum_mans_subarb);

    return ret;
}

int main() {
    fstream f1, f2;
    int x, y;
    char code;

    f1.open("rafaela.in", ios::in);
    f2.open("rafaela.out", ios::out);

    f1 >> N >> Q;

    for(int i=1; i<N; i++) {
        f1 >> x >> y;

        G[x].push_back(y);
        G[y].push_back(x);
    }

    for(int i=1; i<=N; i++) {
        f1 >> initial[i];

        no_of_mans += initial[i];
    }

    create_heavy_path(1);

    for(int i=0; i<Q; i++) {
        f1 >> code;

        if(code == 'Q') {
            f1 >> x;

            f2 << query(x) << '\n';
        }
        else if(code == 'U') {
            f1 >> x;
            f1 >> y;
            no_of_mans += x;
            initial[y] += x;
            update(y, +x);
        }
    }

    f1.close();
    f2.close();

    return 0;
}