#include #include "segment_tree.h" using namespace std; class HeavyLight { public: vector> tree; bool valuesOnVertices; // true - values on vertices, false - values on edges segtree segment_tree; vector parent; vector depth; vector pathRoot; vector in; HeavyLight(const vector> &t, bool valuesOnVertices) : tree(t), valuesOnVertices(valuesOnVertices), segment_tree(t.size()), parent(t.size()), depth(t.size()), pathRoot(t.size()), in(t.size()) { int time = 0; parent[0] = -1; function dfs1 = [&](int u) { int size = 1; int maxSubtree = 0; for (int &v : tree[u]) { if (v == parent[u]) continue; parent[v] = u; depth[v] = depth[u] + 1; int subtree = dfs1(v); if (maxSubtree < subtree) { maxSubtree = subtree; swap(v, tree[u][0]); } size += subtree; } return size; }; function dfs2 = [&](int u) { in[u] = time++; for (int v : t[u]) { if (v == parent[u]) continue; pathRoot[v] = v == t[u][0] ? pathRoot[u] : v; dfs2(v); } }; dfs1(0); dfs2(0); } segtree::node get(int u, int v) { segtree::node res; process_path(u, v, [this, &res](int a, int b) { res = segtree::unite(res, segment_tree.get(a, b)); }); return res; } void modify(int u, int v, long long delta) { process_path(u, v, [this, delta](int a, int b) { segment_tree.modify(a, b, delta); }); } void process_path(int u, int v, const function &op) { for (; pathRoot[u] != pathRoot[v]; v = parent[pathRoot[v]]) { if (depth[pathRoot[u]] > depth[pathRoot[v]]) swap(u, v); op(in[pathRoot[v]], in[v]); } if (u != v || valuesOnVertices) op(min(in[u], in[v]) + (valuesOnVertices ? 0 : 1), max(in[u], in[v])); } }; // usage example int main() { vector> tree{{1, 2}, {0, 3, 4}, {0}, {1}, {1}}; HeavyLight hl_v(tree, true); hl_v.modify(3, 2, 1); hl_v.modify(1, 0, -1); cout << hl_v.get(4, 2).sum << endl; HeavyLight hl_e(tree, false); hl_e.modify(3, 2, 1); hl_e.modify(1, 0, -1); cout << hl_e.get(4, 2).sum << endl; }