Improve HLD

Signed-off-by: Bruno Freitas Tissei <bft15@inf.ufpr.br>

algorithms/graph/hld.cpp

@@ -5,11 +5,14 @@
 ///   - query: O(log^2 n)
 ///   - update: O(log n)
 /// Space: O(n log n)
-
-#define LOG 20 // log2(MAX)
+/// 
+/// Caution:
+///   - Modifications are necessary if values are associated to vertices;
+///   - Modifications are necerrary if other data structure is used.
+/// 
+/// Status: Tested (QTREE,URI2887)
 
 ii edge[MAX];
-int par[MAX][LOG];
 vector<ii> graph[MAX];
  
 template <typename ST>
@@ -17,124 +20,85 @@ struct HLD {
   ST &seg;
   int cnum, ptr;
   
-  // Size of subtree, and depth
-  vector<int> size, dep;
+  // depth, parent, value of node
+  vector<int> dep, par, val;
 
-  // Chain's head, chain's index, base array for segtree, 
-  // position of node in base, and bottommost node
-  vector<int> chead, cidx, base, pos, bot;
+  // head[i]: head of i-th node's chain;
+  // heavy[i]: "special child" of i-th node
+  // pos[i]: position of i-th node on segtree
+  // bot[i]: bottommost (depth-wise) node on the i-th 
+  //   edge (required only when edges have weights)
+  vector<int> head, heavy, pos, bot;
  
   HLD(int n, ST &seg) :
-    seg(seg), size(n), dep(n, 0), bot(n),
-    chead(n, -1), cidx(n), base(n), pos(n)
+    seg(seg), 
+    dep(n, 0), par(n), val(n),
+    head(n), heavy(n, -1), pos(n), bot(n)
   {
-    mset(par, -1);
     cnum = ptr = 0;
  
+    N = n; // global N for segtree
     dfs(0);
     decompose(0);
- 
-    N = ptr; // global N for segtree
-    seg.build(base);
- 
-    // bot[i] stores the bottommost (depth-wise) node
-    // on the i-th edge
-    for (int i = 0; i < N - 1; ++i)
+
+    // (required only when edges have weights)
+    for (int i = 0; i < n - 1; ++i)
       if (dep[edge[i].fi] > dep[edge[i].se])
         bot[i] = edge[i].fi;
       else
         bot[i] = edge[i].se;
   }
  
-  void dfs(int x, int p = -1) {
-    size[x] = 1;
-    par[x][0] = p;
- 
-    for (int i = 1; i < LOG; ++i)
-      if (par[x][i - 1] != -1)
-        par[x][i] = par[par[x][i - 1]][i - 1];
+  int dfs(int x, int p = -1) {
+    int size = 1;
+    par[x] = p;
  
+    int max_size = 0;
     for (auto i : graph[x])
       if (i.fi != p) {
         dep[i.fi] = dep[x] + 1;
-        dfs(i.fi, x);
-        size[x] += size[i.fi];
+        val[i.fi] = i.se;
+        int isize = dfs(i.fi, x);
+
+        size += isize;
+        if (isize > max_size)
+          max_size = isize, heavy[x] = i.fi;
       }
+
+    return size;
   }
  
-  void decompose(int x, int c = -1, int p = -1) {
-    if (chead[cnum] == -1)
-      chead[cnum] = x;
- 
-    pos[x] = ptr;
-    cidx[x] = cnum;
-    base[ptr++] = c;
- 
-    ii sc(-1, -1);
-    if (graph[x].size() > 1) {
-      for (auto i : graph[x])
-        if (i.fi != p)
-          if (sc.fi == -1 || size[sc.fi] < size[i.fi])
-            sc = i;
+  void decompose(int x, int h = 0) {
+    head[x] = h;
+    seg.update(ptr, val[x]);
+    pos[x] = ptr++;
+
+    if (heavy[x] != -1)
+      decompose(heavy[x], h);
 
-      decompose(sc.fi, sc.se, x);
-    }
- 
     for (auto i : graph[x])
-      if (sc.fi != i.fi && i.fi != p) {
-        cnum++;
-        decompose(i.fi, i.se, x); 
-      }
+      if (i.fi != par[x] && i.fi != heavy[x])
+        decompose(i.fi, i.fi);
   }
  
-  int query_up(int a, int b) {
-    if (a == b) return 0;
- 
-    int ca, cb = cidx[b], ans = -1;
-    while (1) {
-      ca = cidx[a];
-      if (ca == cb) {
-        if (a == b) break;
-        ans = max(ans, seg.query(pos[b] + 1, pos[a]));
-        break;
-      }
- 
-      ans = max(ans, seg.query(pos[chead[ca]], pos[a]));
-      a = par[chead[ca]][0];
+  // Queries max edge (or vertice) on the path 
+  // between a and b
+  int query(int a, int b) {
+    int ans = seg.ident;
+    for (; head[a] != head[b]; b = par[head[b]]) {
+      if (dep[head[a]] > dep[head[b]])
+        swap(a, b);
+      ans = seg.op(ans, seg.query(pos[head[b]], pos[b]));
     }
- 
-    return ans;
-  }
- 
-  int lca(int a, int b) {
-    int ans = 0;
- 
-    if (dep[a] < dep[b])
+
+    if (dep[a] > dep[b])
       swap(a, b);
- 
-    for (int i = LOG - 1; i >= 0; --i)
-      if (par[a][i] != -1 && dep[par[a][i]] >= dep[b])
-        a = par[a][i];
- 
-    if (a == b)
-      return a;
- 
-    for (int i = LOG - 1; i >= 0; --i)
-      if (par[a][i] != -1 && par[a][i] != par[b][i]) {
-        a = par[a][i];
-        b = par[b][i];
-      }
- 
-    return par[a][0];
-  }
- 
-  // Queries max edge on the path between a and b
-  int query(int a, int b) {
-    int lab = lca(a, b);
-    return max(query_up(a, lab), query_up(b, lab));
+
+    // Remove "+ 1" when values are associated with vertices
+    return seg.op(ans, seg.query(pos[a] + 1, pos[b]));
   }
  
-  // Updates value of i-th edge
+  // Updates value of i-th edge (or vertice)
   void update(int i, int val) {
     seg.update(pos[bot[i]], val);
   }