[UVA] 13101 - Tobby on Tree

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : 13101 - Tobby on Tree
Source            : UVA Online Judge
Category          : Graph Theory
Algorithm         : Heavy Light Decomposition, Segment Tree
Verdict           : Accepted

#include <bits/stdc++.h>  
  
using namespace std;  
  
  
static const int maxn = 1e5 + 5;  
static const int logn = 18;  
  
int chainNo, chainHead[maxn], chainInd[maxn], chainPos[maxn], chainSize[maxn],  
    baseArray[maxn], posBase[maxn], ptr,  
    depth[maxn], father[maxn][logn], subTreeSize[maxn], starttnode[maxn], endtnode[maxn],  
    S, E, nodeCost[maxn];  
  
vector <int> graph[maxn];  
  
void dfs(int u = 1, int p = -1)  
{  
      for (int i = 1; i < logn; i++) father[u][i] = father[ father[u][i-1] ][i-1];  
      subTreeSize[u] = 1;  
      for (int v : graph[u])  
      {  
            if (v == p) continue;  
            depth[v] = depth[u] + 1;  
            father[v][0] = u;  
            dfs(v, u);  
            subTreeSize[u] += subTreeSize[v];  
      }  
}  
  
int findLCA(int u, int v)  
{  
      if (depth[u] < depth[v]) swap(u, v);  
      for (int i = logn-1; i >= 0; i--)  
      {  
            if (depth[ father[u][i] ] >= depth[v])  
            {  
                  u = father[u][i];  
            }  
      }  
      if (u == v) return u;  
      for (int i = logn-1; i >= 0; i--)  
      {  
            if (father[u][i] != father[v][i])  
            {  
                  u = father[u][i];  
                  v = father[v][i];  
            }  
      }  
      return father[u][0];  
}  
  
void hld(int u = 1, int p = -1)  
{  
      if (chainHead[chainNo] == -1) chainHead[chainNo] = u;  
      chainInd[u] = chainNo;  
      ptr++;  
      baseArray[ptr] = nodeCost[u];  
      posBase[u] = ptr;  
      int bigChild(-1), mx(-1);  
      for (int v : graph[u])  
      {  
            if (v == p) continue;  
            if (subTreeSize[v] > mx) mx = subTreeSize[v], bigChild = v;  
      }  
      if (bigChild != -1) hld(bigChild, u);  
      for (int v : graph[u])  
      {  
            if (v == p) continue;  
            if (bigChild != v)  
            {  
                  chainNo++;  
                  hld(v, u);  
            }  
      }  
}  
  
int Tree[maxn << 2];  
  
void build(int node, int a, int b)  
{  
      if (a == b)  
      {  
            Tree[node] = baseArray[a];  
            return;  
      }  
      int lft = node << 1;  
      int rgt = lft | 1;  
      int mid = (a + b) >> 1;  
      build(lft, a, mid);  
      build(rgt, mid+1, b);  
      Tree[node] = __gcd(Tree[lft], Tree[rgt]);  
}  
  
void update(int node, int a, int b, int pos, int val)  
{  
      if (a > b || a > pos || b < pos) return;  
      if (a >= pos && b <= pos)  
      {  
            Tree[node] = val;  
            return;  
      }  
      int lft = node << 1;  
      int rgt = lft | 1;  
      int mid = (a + b) >> 1;  
      update(lft, a, mid, pos, val);  
      update(rgt, mid+1, b, pos, val);  
      Tree[node] = __gcd(Tree[lft], Tree[rgt]);  
}  
  
int query(int node, int a, int b, int i, int j)  
{  
      if (a > b || a > j || b < i) return 0;  
      if (a >= i && b <= j) return Tree[node];  
      int lft = node << 1;  
      int rgt = lft | 1;  
      int mid = (a + b) >> 1;  
      int p1 = query(lft, a, mid, i, j);  
      int p2 = query(rgt, mid+1, b, i, j);  
      return __gcd(p1, p2);  
}  
  
int query_up(int u, int v)  
{  
      // u niche, v upore  
      int uChain, vChain = chainInd[v];  
      int g = 0;  
      while (true)  
      {  
            uChain = chainInd[u];  
            if (uChain == vChain)  
            {  
                  int l = posBase[v];  
                  int r = posBase[u];  
                  if (l > r) break;  
                  int ans = query(1, S, E, l, r);  
                  g = __gcd(g, ans);  
                  break;  
            }  
            int l = posBase[ chainHead[uChain] ];  
            int r = posBase[u];  
            int ans = query(1, S, E, l, r);  
            g = __gcd(g, ans);  
            u = father[ chainHead[uChain] ][0];  
      }  
      return g;  
}  
  
int query_hld(int u, int v)  
{  
      int lca = findLCA(u, v);  
      int ans = 0;  
      ans = __gcd(ans, query_up(u, lca));  
      ans = __gcd(ans, query_up(v, lca));  
      return ans;  
}  
  
void change_node_cost(int u, int val)  
{  
      int pos = posBase[u];  
      nodeCost[u] = val;  
      update(1, S, E, pos, val);  
}  
  
void init()  
{  
      ptr = 0;  
      chainNo = 0;  
      for (int i = 0; i < maxn; i++)  
      {  
            graph[i].clear();  
            subTreeSize[i] = 0;  
            chainHead[i] = -1;  
            depth[i] = 0;  
            baseArray[i] = 0;  
            posBase[i] = 0;  
            fill(begin(father[i]), end(father[i]), 0);  
      }  
}  
  
int main()  
{ 
      int tnode;  
      while (scanf("%d", &tnode) == 1)  
      {  
            init();  
            for (int i = 1; i <= tnode; i++) scanf("%d", nodeCost+i);  
            for (int e = 1; e < tnode; e++)  
            {  
                  int u, v;  
                  scanf("%d %d", &u, &v);  
                  u++; v++;  
                  graph[u].push_back(v);  
                  graph[v].push_back(u);  
            }  
            depth[1] = 1;  
            dfs();  
            hld();  
            S = 1;  
            E = ptr;  
            build(1, S, E);  
            int tquery;  
            scanf("%d", &tquery);  
            while (tquery--)  
            {  
                  int type;  
                  scanf("%d", &type);  
                  if (type == 1)  
                  {  
                        int u, v;  
                        scanf("%d %d", &u, &v);  
                        u++; v++;  
                        int ans = query_hld(u, v);  
                        printf("%d\n", ans);  
                  }  
                  else  
                  {  
                        int u, x;  
                        scanf("%d %d", &u, &x);  
                        u++;  
                        change_node_cost(u, x);  
                  }  
            }  
      }  
}