Saturday, February 2, 2019

[UVA] 12424 - Answering Queries on a Tree

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

  1. #include <bits/stdc++.h>  
  2.   
  3. using namespace std;  
  4.   
  5. static const int maxn = 1e5 + 5;  
  6. static const int logn = 19;  
  7.   
  8. int chainNo, chainHead[maxn], chainInd[maxn], chainPos[maxn], chainSize[maxn],  
  9.     baseArray[maxn], posBase[maxn], ptr, depth[maxn], father[maxn][logn],  
  10.     subTreeSize[maxn], startnode[maxn], endNode[maxn], nodeInd[maxn],  
  11.     S, E, nodeColor[maxn];  // color of an edge stored in a node wrt dfs tree  
  12.   
  13. struct node  
  14. {  
  15.       int v,  // child node  
  16.           e;  // edge number  
  17.       node() {}  
  18.       node(int v, int e) : v(v), e(e) {}  
  19. };  
  20.   
  21. vector <node> graph[maxn];  
  22.   
  23. void dfs(int u = 1, int p = -1)  
  24. {  
  25.       for (int i = 1; i < logn; i++) father[u][i] = father[father[u][i-1]][i-1];  
  26.       subTreeSize[u] = 1;  
  27.       for (node g : graph[u])  
  28.       {  
  29.             int v = g.v;  
  30.             int e = g.e;  
  31.             if (v == p) continue;  
  32.             depth[v] = depth[u] + 1;  
  33.             father[v][0] = u;  
  34.             startnode[e] = u;  
  35.             endNode[e] = v;  
  36.             dfs(v, u);  
  37.             subTreeSize[u] += subTreeSize[v];  
  38.       }  
  39. }  
  40.   
  41. int findLCA(int u, int v)  
  42. {  
  43.       if (depth[u] < depth[v]) swap(u, v);  
  44.       for (int i = logn-1; i >= 0; i--)  
  45.       {  
  46.             if (depth[father[u][i]] >= depth[v])  
  47.             {  
  48.                   u = father[u][i];  
  49.             }  
  50.       }  
  51.       if (u == v) return u;  
  52.       for (int i = logn-1; i >= 0; i--)  
  53.       {  
  54.             if (father[u][i] != father[v][i])  
  55.             {  
  56.                   u = father[u][i];  
  57.                   v = father[v][i];  
  58.             }  
  59.       }  
  60.       return father[u][0];  
  61. }  
  62.   
  63. void hld(int u = 1, int p = -1)  
  64. {  
  65.       if (chainHead[chainNo] == -1) chainHead[chainNo] = u;  
  66.       chainInd[u] = chainNo;  
  67.       ptr++;  
  68.       baseArray[ptr] = nodeColor[u];  
  69.       posBase[u] = ptr;  
  70.       int bigChild = -1;  
  71.       int mx = -1;  
  72.       for (node g : graph[u])  
  73.       {  
  74.             int v = g.v;  
  75.             if (v == p) continue;  
  76.             if (subTreeSize[v] > mx) mx = subTreeSize[v], bigChild = v;  
  77.       }  
  78.       if (bigChild != -1) hld(bigChild, u);  
  79.       for (node g : graph[u])  
  80.       {  
  81.             int v = g.v;  
  82.             if (v == p) continue;  
  83.             if (bigChild != v)  
  84.             {  
  85.                   chainNo++;  
  86.                   hld(v, u);  
  87.             }  
  88.       }  
  89. }  
  90.   
  91. struct segmentTree  
  92. {  
  93.       int color[11];  // store number of color  
  94.       void assignLeaf(int val)  
  95.       {  
  96.             fill(begin(color), end(color), 0);  
  97.             assert(1 <= val && val <= 10);  
  98.             color[val]++;  
  99.       }  
  100.       void marge(segmentTree &L, segmentTree &R)  
  101.       {  
  102.             for (int i = 1; i <= 10; i++) color[i] = L.color[i] + R.color[i];  
  103.       }  
  104.       int maxTimes()  
  105.       {  
  106.             int maxc = -1;  
  107.             for (int i = 1; i <= 10; i++) maxc = max(maxc, color[i]);  
  108.       }  
  109. } Tree[maxn << 2];  
  110.   
  111. void build(int node, int a, int b)  
  112. {  
  113.       if (a == b)  
  114.       {  
  115.             Tree[node].assignLeaf(baseArray[a]);  
  116.             return;  
  117.       }  
  118.       int lft = node << 1;  
  119.       int rgt = lft | 1;  
  120.       int mid = (a + b) >> 1;  
  121.       build(lft, a, mid);  
  122.       build(rgt, mid+1, b);  
  123.       Tree[node].marge(Tree[lft], Tree[rgt]);  
  124. }  
  125.   
  126. void update(int node, int a, int b, int pos, int value)  
  127. {  
  128.       if (a > b || a > pos || b < pos) return;  
  129.       if (a >= pos && b <= pos)  
  130.       {  
  131.             Tree[node].assignLeaf(value);  
  132.             return;  
  133.       }  
  134.       int lft = node << 1;  
  135.       int rgt = lft | 1;  
  136.       int mid = (a + b) >> 1;  
  137.       update(lft, a, mid, pos, value);  
  138.       update(rgt, mid+1, b, pos, value);  
  139.       Tree[node].marge(Tree[lft], Tree[rgt]);  
  140. }  
  141.   
  142. segmentTree query(int node, int a, int b, int i, int j)  
  143. {  
  144.       if (a > b || a > j || b < i)  
  145.       {  
  146.             segmentTree nul;  
  147.             fill(begin(nul.color), end(nul.color), 0);  
  148.             return nul;  
  149.       }  
  150.       if (a >= i && b <= j) return Tree[node];  
  151.       int lft = node << 1;  
  152.       int rgt = lft | 1;  
  153.       int mid = (a + b) >> 1;  
  154.       segmentTree p1, p2, res;  
  155.       p1 = query(lft, a, mid, i, j);  
  156.       p2 = query(rgt, mid+1, b, i, j);  
  157.       res.marge(p1, p2);  
  158.       return res;  
  159. }  
  160.   
  161. void change_node_color(int node, int value)  
  162. {  
  163.       int pos = posBase[node];  
  164.       baseArray[pos] = value;  
  165.       update(1, S, E, pos, value);  
  166. }  
  167.   
  168. int qColor[11];  
  169.   
  170. void query_UP(int u, int v)  
  171. {  
  172.       int uChian,  
  173.           vChain = chainInd[v];  
  174.       while (true)  
  175.       {  
  176.             uChian = chainInd[u];  
  177.             if (uChian == vChain)  
  178.             {  
  179.                   int l = posBase[v];  
  180.                   int r = posBase[u];  
  181.                   if (l > r) break;  
  182.                   segmentTree q = query(1, S, E, l, r);  
  183.                   for (int i = 1; i <= 10; i++) qColor[i] += q.color[i];  
  184.                   break;  
  185.             }  
  186.             int l = posBase[chainHead[uChian]];  
  187.             int r = posBase[u];  
  188.             segmentTree q = query(1, S, E, l, r);  
  189.             for (int i = 1; i <= 10; i++) qColor[i] += q.color[i];  
  190.             u = father[chainHead[uChian]][0];  
  191.       }  
  192. }  
  193.   
  194. int query_hld(int u, int v)  
  195. {  
  196.       if (u == v) return 1;  
  197.       int lca = findLCA(u, v);  
  198.       fill(begin(qColor), end(qColor), 0);  
  199.       if (u != lca) query_UP(u, lca);  
  200.       if (v != lca) query_UP(v, lca);  
  201.       if (u != lca && v != lca)  // we calculate this node twice  
  202.       {  
  203.             qColor[baseArray[posBase[lca]]]--;  
  204.       }  
  205.       int maxc = 0;  
  206.       for (int i = 1; i <= 10; i++)  
  207.       {  
  208.             if (qColor[i] > maxc) maxc = qColor[i];  
  209.       }  
  210.       return maxc;  
  211. }  
  212.   
  213.   
  214. void initialize()  
  215. {  
  216.       ptr = 0;  
  217.       chainNo = 0;  
  218.       for (int i = 0; i < maxn; i++)  
  219.       {  
  220.             graph[i].clear();  
  221.             subTreeSize[i] = 0;  
  222.             chainHead[i] = -1;  
  223.             chainInd[i] = 0;  
  224.             depth[i] = 0;  
  225.             baseArray[i] = 0;  
  226.             posBase[i] = 0;  
  227.             for (int j = 0; j < logn; j++) father[i][j] = 0;  
  228.       }  
  229. }  
  230.   
  231. int main()  
  232. {  
  233.       //freopen("in.txt", "r", stdin);  
  234.   
  235.       int tc;  
  236.       scanf("%d", &tc);  
  237.       for (int tcase = 1; tcase <= tc; tcase++)  
  238.       {  
  239.             initialize();  
  240.             int tnode, tquery;  
  241.             scanf("%d %d", &tnode, &tquery);  
  242.             for (int i = 1; i <= tnode; i++) scanf("%d", nodeColor+i);  
  243.             for (int e = 1; e < tnode; e++)  
  244.             {  
  245.                   int u, v;  
  246.                   scanf("%d %d", &u, &v);  
  247.                   graph[u].push_back({v, e});  
  248.                   graph[v].push_back({u, e});  
  249.             }  
  250.             int root = 1;  
  251.             depth[root] = 1;  
  252.             dfs();  
  253.             hld();  
  254.             S = 1;  
  255.             E = ptr;  
  256.             build(1, S, E);  
  257.             while (tquery--)  
  258.             {  
  259.                   int type, u, v;  
  260.                   scanf("%d %d %d", &type, &u, &v);  
  261.                   if (type == 0)  
  262.                   {  
  263.                         change_node_color(u, v);  
  264.                   }  
  265.                   else  
  266.                   {  
  267.                         int maxc = query_hld(u, v);  
  268.                         printf("%d\n", maxc);  
  269.                   }  
  270.             }  
  271.       }  
  272. }  

No comments:

Post a Comment

Note: Only a member of this blog may post a comment.