From 89e07ebc8aeb9ee95801cf420f66621c97ebe28c Mon Sep 17 00:00:00 2001
From: subcrip <contact@subc.rip>
Date: Wed, 22 May 2024 11:48:26 +0800
Subject: [PATCH] Update graph/hld.cc

Signed-off-by: subcrip <contact@subc.rip>
---
 graph/hld.cc | 65 +++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 59 insertions(+), 6 deletions(-)

diff --git a/graph/hld.cc b/graph/hld.cc
index d0c6db6..fec48e6 100644
--- a/graph/hld.cc
+++ b/graph/hld.cc
@@ -1,13 +1,18 @@
-namespace HLD {
+template <typename Info, typename Tag>
+struct HLD {
     struct node_info {
         int father, depth, hson, size, head, dfn = -1;
     };
+
+    int n;
+    vector<int> seq;
+    vector<node_info> info;
+    segtree<Info, Tag> tr;
+
     // returns: (dfs sequence, node info)
     // node numbering starts from `1`
     // if `dfn(v) == -1`, then node `v` is never accessed.
-    pair<vector<int>, vector<node_info>> work(const vector<vector<int>>& ch, int root = 0) {
-        int n = ch.size() - 1;
-        vector<int> seq;
+    HLD(const vector<vector<int>>& ch, const vector<Info>& init, int root = 0) : n(ch.size() - 1), seq(), info(n + 1), tr(n + 1) {
         vector<node_info> res(n + 1);
         auto dfs1 = [&] (auto dfs1, int v, int pa) -> void {
             res[v].father = pa;
@@ -38,6 +43,54 @@ namespace HLD {
             }
         };
         dfs2(dfs2, root, root);
-        return { seq, res };
+        info = res;
+
+        for (int i = 1; i <= n; ++i) {
+            tr.set(info[i].dfn, init[i]);
+        }
     }
-}
+
+    void set(int v, const Info& t) {
+        tr.set(info[v].dfn, t);
+    }
+
+    void apply(int v, const Tag& t) {
+        tr.apply(info[v].dfn, t);
+    }
+
+    Info query(int v) {
+        return tr.query(info[v].dfn);
+    }
+
+    void path_apply(int u, int v, const Tag& t) {
+        while (info[u].head != info[v].head) {
+            if (info[info[u].head].depth < info[info[v].head].depth) {
+                swap(u, v);
+            }
+            tr.range_apply(info[info[u].head].dfn, info[u].dfn, t);
+        }
+        if (info[u].depth < info[v].depth) swap(u, v);
+        tr.range_apply(info[v].dfn, info[u].dfn);
+    }
+
+    Info path_query(int u, int v) {
+        Info res;
+        while (info[u].head != info[v].head) {
+            if (info[info[u].head].depth < info[info[v].head].depth) {
+                swap(u, v);
+            }
+            res = res + tr.range_query(info[info[u].head].dfn, info[u].dfn);
+        }
+        if (info[u].depth < info[v].depth) swap(u, v);
+        res = res + tr.range_query(info[v].dfn, info[u].dfn);
+        return res;
+    }
+
+    void subtree_apply(int v, const Tag& t) {
+        tr.range_apply(info[v].dfn, info[v].dfn + info[v].size - 1, t);
+    }
+
+    Info subtree_query(int v) {
+        return tr.range_query(info[v].dfn, info[v].dfn + info[v].size - 1);
+    }
+};