HLD.cpp

/*
    Author : M. A. Rafsan Mazumder
    Algorithm : Heavy Light Decomposition
    Problem : SPOJ - QTREE
              Given a tree with N nodes, edges are numbered 1, 2, 3, .. , n-1
              1-Query : CHANGE i ti : Change the cost of the ith edge to ti
              2-Query : a b : ask for the maximum edge cost on the path from node a to node b
*/

#include <bits/stdc++.h>

using namespace std;

#define MAX 10005
#define MAX_LOG 15

vector<int> adj[MAX], costs[MAX];
int n, root, subSize[MAX], par[MAX][MAX_LOG], level[MAX];
int baseArray[MAX], ptr;
int chainNo, chainInd[MAX], chainHead[MAX], posInBase[MAX];
int chainSize[MAX], chainPos[MAX];
int tree[4*MAX];
pair<pair<int, int>, int> p[MAX];

void gen()
{
    for(int i=0; i<MAX; i++){
        adj[i].clear();
        costs[i].clear();
        chainHead[i] = -1;
        chainSize[i] = 0;
        for(int j=0; j<MAX_LOG; j++) par[i][j] = -1;
    }
    ptr = 1;
    chainNo = 1;
    root = 1;
}

/*
 * dfs used to set parent of a node, depth of a node, subtree size of a node
 */

void dfs(int u, int prev, int depth)
{
    par[u][0] = prev;
    level[u] = depth;
    subSize[u] = 1;

    for(int i=0; i<adj[u].size(); i++){
        int v = adj[u][i];
        if(v != prev){
            dfs(v, u, depth+1);
            subSize[u] += subSize[v];
        }
    }
}

/*
 * precompute 2^j the parent of ith node
 */

void preComputeLCA()
{
    for(int j=1; j<MAX_LOG; j++){
        for(int i=1; i<=n; i++){
            if(par[i][j-1] != -1) par[i][j] = par[par[i][j-1]][j-1];
        }
    }
}

/*
 * LCA: Takes two nodes u, v and returns Lowest Common Ancestor of u, v
 */

int LCA(int u, int v)
{
    if(level[u] < level[v]) swap(u, v);
    int diff = level[u]-level[v];
    for(int i=0; i<MAX_LOG; i++){
        if((1<<i) & diff) u = par[u][i];
    }
    if(u == v) return u;
    for(int i=MAX_LOG-1; i>=0; i--){
        if(par[u][i] != par[v][i]){
            u = par[u][i];
            v = par[v][i];
        }
    }
    return par[u][0];
}

/*
 * Actual HL-Decomposition part
 * Initially all entries of chainHead[] are set to -1.
 * So when ever a new chain is started, chain head is correctly assigned.
 * As we add a new node to chain, we will note its position in the baseArray.
 * In the first for loop we find the child node which has maximum sub-tree size.
 * The following if condition is failed for leaf nodes.
 * When the if condition passes, we expand the chain to special child.
 * In the second for loop we recursively call the function on all normal nodes.
 * chainNo++ ensures that we are creating a new chain for each normal child.
 */

 /*
  * chainInd[curNode] - Given a node, to which chain does that node belong to.
  * chainHead[chainNo] - Given a chain, what is the head of the chain
  * chainSize[chainNo] - Given a chain, what is the size of the chain
  * chainPos[curNode] - Given a node, what is the position of the node in its chain
  * posInBase[curNode] -  Position of this node in baseArray which we will use in Segtree
  */

void HLD(int curNode, int cost, int prev)
{
    if(chainHead[chainNo] == -1){
        chainHead[chainNo] = curNode; // Assign chain head
    }
    chainInd[curNode] = chainNo;
    chainSize[chainNo]++;
    chainPos[curNode] = chainSize[chainNo]; // 1-indexed
    posInBase[curNode] = ptr; // Position of this node in baseArray which we will use in Segtree
    baseArray[ptr] = cost;
    ptr++;

    int sc = -1, ncost;
    // Loop to find special child
    for(int i=0; i<adj[curNode].size(); i++){
        int v = adj[curNode][i];
        if(v != prev){
            if(sc == -1 || subSize[v] > subSize[sc]){
                sc = v;
                ncost = costs[curNode][i];
            }
        }
    }

    if(sc != -1){
        // Expand the chain
        HLD(sc, ncost, curNode);
    }

    for(int i=0; i<adj[curNode].size(); i++){
        int v = adj[curNode][i];
        if(v != prev && v != sc){
            // New chains at each normal node
            chainNo++;
            HLD(v, costs[curNode][i], curNode);
        }
    }
}

/*
 * init_tree:
 * Used to construct the segment tree. It uses the baseArray for construction
 */

void init_tree(int node, int b, int e)
{
    if(b == e){
        tree[node] = baseArray[b];
        return;
    }
    int left = 2*node;
    int right = 2*node+1;
    int mid = (b+e)/2;
    init_tree(left, b, mid);
    init_tree(right, mid+1, e);
    if(tree[left] > tree[right]) tree[node] = tree[left];
    else tree[node] = tree[right];
}

/*
 * query_tree:
 * Given S and E, it will return the maximum value in the range [S,E]
 */

int query_tree(int node, int b, int e, int l, int r)
{
    if(l > e || r < b) return 0;
    if(b >= l && e <= r) return tree[node];
    int left = 2*node;
    int right = 2*node+1;
    int mid = (b+e)/2;
    int q1 = query_tree(left, b, mid, l, r);
    int q2 = query_tree(right, mid+1, e, l, r);
    if(q1 > q2) return q1;
    else return q2;
}

/*
 * update_tree:
 * Point update. Update a single element of the segment tree.
 */

void update_tree(int node, int b, int e, int pos, int val)
{
    if(pos > e || pos < b) return;
    if(b == pos && e == pos){
        tree[node] = val;
        return;
    }
    int left = 2*node;
    int right = 2*node+1;
    int mid = (b+e)/2;
    update_tree(left, b, mid, pos, val);
    update_tree(right, mid+1, e, pos, val);
    if(tree[left] > tree[right]) tree[node] = tree[left];
    else tree[node] = tree[right];
}

/*
 * query_up:
 * It takes two nodes u and v, condition is that v is an ancestor of u
 * We query the chain in which u is present till chain head, then move to next chain up
 * We do that way till u and v are in the same chain, we query for that part of chain and break
 */

 int query_up(int u, int v)
 {
     if(u == v) return 0;
     int uchain, vchain = chainInd[v], ans = -1;
     while(1){
        uchain = chainInd[u];
        if(uchain == vchain){
            // Both u and v are in the same chain, so we need to query from u to v, update answer and break.
            // We break because we came from u up till v, we are done
            if(u == v) break;
            int temp = query_tree(1, 1, ptr, posInBase[v]+1, posInBase[u]);
            if(temp > ans) ans = temp;
            break;
        }
        int temp = query_tree(1, 1, ptr, posInBase[chainHead[uchain]], posInBase[u]);
        if(temp > ans) ans = temp;
        u = chainHead[uchain]; // move u to u's chainHead
        u = par[u][0]; // Then move to its parent, that means we changed chains
     }
     return ans;
 }


/*
 * We have a query from u to v, we break it into two queries, u to LCA(u, v) and
                                                              v to LCA(u, v)
 */

void query(int u, int v)
{
    int lca = LCA(u, v);
    int ans1 = query_up(u, lca);
    int ans2 = query_up(v, lca);
    int ans;
    if(ans1 > ans2) ans = ans1;
    else ans = ans2;
    printf("%d\n", ans);
}

void change(int edgeNum, int val)
{
    int u = p[edgeNum].first.first;
    int v = p[edgeNum].first.second;
    int cost = p[edgeNum].second;
    int node = u;
    if(level[v] > level[u]) node = v;
    update_tree(1, 1, ptr, posInBase[node], val);
}

int main()
{
    //freopen("in.txt", "r", stdin);
    int cases;
    scanf("%d", &cases);
    while(cases--){
        gen();
        scanf("%d", &n);
        for(int i=1; i<n; i++){
            int u, v, w;
            scanf("%d %d %d", &u, &v, &w);
            p[i] = make_pair(make_pair(u, v), w);
            costs[u].push_back(w);
            costs[v].push_back(w);
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
        dfs(1, -1, 0);
        preComputeLCA();
        HLD(1, -1, -1);
        ptr--;
        init_tree(1, 1, ptr);
        while(1){
            char s[100];
            scanf("%s", s);
            if(s[0] == 'D') break;
            int a, b;
            scanf("%d %d", &a, &b);
            if(s[0] == 'Q'){
                query(a, b);
            }
            else {
                change(a, b);
            }
        }
    }
}

/*
    1

    3
    1 2 1
    2 3 2
    QUERY 1 2
    CHANGE 1 3
    QUERY 1 2
    DONE

    1
    3
*/