QTREE Solution

//
// QTREE.cpp
//
// SPOJ 375: http://www.spoj.com/problems/QTREE/
// 
// Siwakorn Srisakaokul - ping128
// Written on Tuesday,  3 June 2014.
//

#include <cstdio>
#include <iostream>
#include <sstream>
#include <cstdlib>
#include <string>
#include <vector>
#include <set>
#include <queue>
#include <stack>
#include <list>
#include <cmath>
#include <algorithm>
#include <map>
#include <ctype.h>
#include <string.h>

#include <assert.h>

#define x first
#define y second

using namespace std;

template<class T1, class T2>
ostream& operator << (ostream &out, pair<T1, T2> pair) { return out << "(" << pair.first << ", " << pair.second << ")"; }

typedef long long LL;
typedef pair<int, int> PII;
typedef pair<PII, int> PII2;

#define MAXN 10005
#define MAXLVL 14 // about log(MAXN)

int N;

int edges[MAXN][3]; // u, v, c

vector<PII> adj[MAXN]; // v, c

// Tree info
int depths[MAXN], subtree_sz[MAXN];
int parents[MAXLVL][MAXN]; // parents[0][i] = the parent of node i

// Heavy-Light Decomposition
int num_chain;
int chain_head[MAXN];
int chain_index[MAXN];
int chain_pos[MAXN]; // not neccessary
int chain_size[MAXN]; // not neccessary

int maxx[MAXN];

// lining all chains in the same array to build a segment tree over it
int base_array_ptr;
int pos_in_base_array[MAXN];
int base_array_value[MAXN];  // [ -, chain1's elements, chain2's elements, ... ]

// Set up subtree_sz, depth, parent for each node
void tree_setup(int at, int par, int dep) {
    depths[at] = dep;
    subtree_sz[at] = 1;
    parents[0][at] = par;
    for (int j = 1; j < MAXLVL; j++) parents[j][at] = -1;
    for (int i = 0; i < adj[at].size(); i++) {
        int v = adj[at][i].first;
        if (v != par) {
            tree_setup(v, at, dep + 1);
            subtree_sz[at] += subtree_sz[v];
        }
    }
}

///////////////////////////////
// LCA
///////////////////////////////

// Requires tree_setup()
void lca_setup() {
    for (int i = 1; i < MAXLVL; i++) {
        for (int j = 0; j < N; j++) {
            if (parents[i - 1][j] != -1) {
                parents[i][j] = parents[i - 1][parents[i - 1][j]];
            }
        }
    }
}

int lca(int u, int v) {
    if (depths[u] < depths[v]) swap(u, v);
    int diff_dep = depths[u] - depths[v];
    for (int i = 0; i < MAXLVL; i++)
        if ((1<<i) & diff_dep)
            u = parents[i][u];
    if (u == v) return u;
    for (int i = MAXLVL - 1; i >= 0; i--) {
        if (parents[i][u] != parents[i][v]) {
            u = parents[i][u];
            v = parents[i][v];
        }
    }
    return parents[0][u];
}
            
/////////////////////
// Segment Tree
/////////////////////

typedef struct segment_st {
    int value;
    int left, right;
}SegmentTree;

SegmentTree tree[MAXN * 4];

void updateUp(int at){
    // Recalculate the parent's value without considering the parent's lazy
    tree[at].value = max(tree[at * 2].value, tree[at * 2 + 1].value);
}

void build_tree(int at, int ll, int rr){
    tree[at].left = ll;
    tree[at].right = rr;
    tree[at].value = 0;
    if(ll == rr){
        // Set to some initial value
        tree[at].value = base_array_value[ll];
    } else {
        int mid = (tree[at].left + tree[at].right) / 2;
        build_tree(at * 2, tree[at].left, mid);
        build_tree(at * 2 + 1, mid + 1, tree[at].right);
        updateUp(at);
    }
}

int query(int at, int ll, int rr){
    if(tree[at].left == ll && tree[at].right == rr){
        // Return the SegmentTree[at] value
        return tree[at].value;
    } else {
        int res = 0;
        int mid = (tree[at].left + tree[at].right) / 2;
        if(ll <= mid) res = max(res, query(at * 2, ll, min(rr, mid)));
        if(rr > mid) res = max(res, query(at * 2 + 1, max(mid + 1, ll), rr));
        return res;
    }
}

void update(int at, int ll, int rr, int value){
    if(tree[at].left == ll && tree[at].right == rr){
        tree[at].value = value;
    } else {
        int mid = (tree[at].left + tree[at].right) / 2;
        if(ll <= mid) update(at * 2, ll, min(rr, mid), value);
        if(rr > mid) update(at * 2 + 1, max(mid + 1, ll), rr, value);
        updateUp(at);
    }
}

///////////////////////////////
// HLD
///////////////////////////////

void HLD(int at, int par) {
    if (chain_head[num_chain - 1] == -1) {
        chain_head[num_chain - 1] = at;
    }
    chain_index[at] = num_chain - 1;
    chain_pos[at] = chain_size[num_chain - 1]++;

    pos_in_base_array[at] = base_array_ptr++;

    int temp = -1;
    // Find the heaviest subtree
    for (int i = 0; i < adj[at].size(); i++) {
        int v = adj[at][i].first;
        if (v != par && (temp == -1 || subtree_sz[temp] < subtree_sz[v])) {
            temp = v;
        }
    }

    if (temp != -1) {
        // Expand the chain
        HLD(temp, at);
    }

    for (int i = 0; i < adj[at].size(); i++) {
        int v = adj[at][i].first;
        if (v != par) {
            if (v != temp) {
                // Create a new chain on each small child
                num_chain++;
                HLD(v, at);
            }
            // Set up base_array value
            base_array_value[pos_in_base_array[v]] = adj[at][i].second;
        }
    }
}

void build_HLD(int root) {
    for (int i = 0; i < MAXN; i++) {
        chain_size[i] = 0;
        chain_head[i] = -1;
    }
    num_chain = 1;
    base_array_ptr = 1;
    HLD(root, -1);    
}

// Query up on HLD, ** u is below v
int query_up(int u, int v) {
    if (u == v) {
        return 0;
    }
    int u_chain, v_chain = chain_index[v], ans = 0;
    while (1) {
        u_chain = chain_index[u];
        if (u_chain == v_chain) {
            ans = max(ans, query(1, pos_in_base_array[v] + 1, pos_in_base_array[u]));
            break;
        }
        int temp = query(1, pos_in_base_array[chain_head[u_chain]], pos_in_base_array[u]);
        ans = max(ans, temp);
        u = chain_head[u_chain];
        u = parents[0][u];
    }
    
    return ans;
}

/* Return the maximum value of nodes along the path from u to v */
int max_uv(int u, int v) {
    int x = lca(u, v);
    return max(query_up(u, x), query_up(v, x));
}

void solve() {

    scanf("%d", &N);

    for (int i = 0; i <= N; i++) {
        adj[i].clear();
    }

    for (int i = 0; i < N - 1; i++) {
        scanf("%d %d %d", &edges[i][0], &edges[i][1], &edges[i][2]);
        edges[i][0]--, edges[i][1]--;
        adj[edges[i][0]].push_back(PII(edges[i][1], edges[i][2]));
        adj[edges[i][1]].push_back(PII(edges[i][0], edges[i][2]));
    }

    tree_setup(0, -1, 0);
    lca_setup();
    build_HLD(0); // Creates chains listing in base_array_value for a segment tree to use
    build_tree(1, 1, N); // Sets up a segment tree or other data structures on top of base_array_value
    
    char s[10];
    while (1) {
        scanf("%s", s);
        if (s[0] == 'D') break;
        if (s[0] == 'Q') {
            int u, v;
            scanf("%d %d", &u, &v);
            u--, v--;
            if (u == v) printf("0\n");
            else printf("%d\n", max_uv(u, v));
        } else {
            int id_e, v;
            scanf("%d %d", &id_e, &v);
            id_e--;
            int id_node = (depths[edges[id_e][0]] > depths[edges[id_e][1]]) ? edges[id_e][0] : edges[id_e][1];
            int ind = pos_in_base_array[id_node];
            update(1, ind, ind, v);
        }
    }
}

int main() {
    int test;
    scanf("%d", &test);
    for (int tt = 0; tt < test; tt++) {
        solve();
    }
    
    return 0;
}

No comments:

Post a Comment