The following C++ implementation demonstrates a complete Heavy-Light Decomposition (HLD) framework integrated with a lazy propagation segment tree to support efficient path updates and queries on trees. It passes the standard template problem on Luogu.
#include <bits/stdc++.h>
using namespace std;
using i64 = long long;
int MOD;
struct SegNode {
int l, r;
i64 sum = 0, add = 0;
};
template<typename T>
struct LazySegTree {
#define lc (u << 1)
#define rc (u << 1 | 1)
int n;
vector<T> tree;
LazySegTree(int size) : n(size), tree(4 * size + 10) {
build(1, 1, n);
}
void build(int u, int l, int r) {
tree[u] = {l, r, 0, 0};
if (l == r) return;
int mid = (l + r) / 2;
build(lc, l, mid);
build(rc, mid + 1, r);
}
void apply(int u, i64 val) {
tree[u].sum = (tree[u].sum + val * (tree[u].r - tree[u].l + 1)) % MOD;
tree[u].add = (tree[u].add + val) % MOD;
}
void push(int u) {
if (tree[u].add) {
apply(lc, tree[u].add);
apply(rc, tree[u].add);
tree[u].add = 0;
}
}
void pull(int u) {
tree[u].sum = (tree[lc].sum + tree[rc].sum) % MOD;
}
void update(int u, int l, int r, i64 val) {
if (l > tree[u].r || r < tree[u].l) return;
if (l <= tree[u].l && tree[u].r <= r) {
apply(u, val);
return;
}
push(u);
update(lc, l, r, val);
update(rc, l, r, val);
pull(u);
}
i64 query(int u, int l, int r) {
if (l > tree[u].r || r < tree[u].l) return 0;
if (l <= tree[u].l && tree[u].r <= r) return tree[u].sum;
push(u);
return (query(lc, l, r) + query(rc, l, r)) % MOD;
}
};
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, q, root;
cin >> n >> q >> root >> MOD;
vector<i64> init_val(n + 1);
for (int i = 1; i <= n; ++i) {
cin >> init_val[i];
init_val[i] %= MOD;
}
vector<vector<int>> adj(n + 1);
for (int i = 1; i < n; ++i) {
int u, v;
cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
// First DFS: compute parent, depth, size, heavy child
vector<int> parent(n + 1), depth(n + 1), heavy(n + 1, -1), subtree_sz(n + 1);
function<void(int, int)> dfs1 = [&](int u, int p) {
parent[u] = p;
depth[u] = depth[p] + 1;
subtree_sz[u] = 1;
int max_size = 0;
for (int v : adj[u]) {
if (v == p) continue;
dfs1(v, u);
subtree_sz[u] += subtree_sz[v];
if (subtree_sz[v] > max_size) {
max_size = subtree_sz[v];
heavy[u] = v;
}
}
};
dfs1(root, 0);
// Second DFS: assign chain top and DFS order
vector<int> top(n + 1), pos(n + 1), rev_pos(n + 1);
int timer = 0;
function<void(int, int)> dfs2 = [&](int u, int chain_top) {
top[u] = chain_top;
pos[u] = ++timer;
rev_pos[timer] = u;
if (heavy[u] != -1) dfs2(heavy[u], chain_top);
for (int v : adj[u]) {
if (v != parent[u] && v != heavy[u]) {
dfs2(v, v);
}
}
};
dfs2(root, root);
// Build segment tree on DFS order
vector<i64> arr(n + 1);
for (int i = 1; i <= n; ++i) arr[pos[i]] = init_val[i];
LazySegTree<SegNode> seg_tree(n);
for (int i = 1; i <= n; ++i) seg_tree.update(1, i, i, arr[i]);
// Path update between u and v
auto path_update = [&](int u, int v, i64 delta) {
while (top[u] != top[v]) {
if (depth[top[u]] < depth[top[v]]) swap(u, v);
seg_tree.update(1, pos[top[u]], pos[u], delta);
u = parent[top[u]];
}
if (depth[u] > depth[v]) swap(u, v);
seg_tree.update(1, pos[u], pos[v], delta);
};
// Path query between u and v
auto path_query = [&](int u, int v) -> i64 {
i64 res = 0;
while (top[u] != top[v]) {
if (depth[top[u]] < depth[top[v]]) swap(u, v);
res = (res + seg_tree.query(1, pos[top[u]], pos[u])) % MOD;
u = parent[top[u]];
}
if (depth[u] > depth[v]) swap(u, v);
res = (res + seg_tree.query(1, pos[u], pos[v])) % MOD;
return res;
};
// Subtree update/query uses [pos[u], pos[u] + subtree_sz[u] - 1]
while (q--) {
int op;
cin >> op;
if (op == 1) {
int x, y, z;
cin >> x >> y >> z;
path_update(x, y, z % MOD);
} else if (op == 2) {
int x, y;
cin >> x >> y;
cout << path_query(x, y) << '\n';
} else if (op == 3) {
int x, z;
cin >> x >> z;
seg_tree.update(1, pos[x], pos[x] + subtree_sz[x] - 1, z % MOD);
} else {
int x;
cin >> x;
cout << seg_tree.query(1, pos[x], pos[x] + subtree_sz[x] - 1) << '\n';
}
}
return 0;
}
Core HLD Components
First DFS (Compute Sizes and Heavy Children):
vector<int> parent(n+1), depth(n+1), heavy(n+1,-1), sz(n+1);
function<void(int,int)> dfs1 = [&](int u, int p) {
parent[u] = p;
depth[u] = depth[p] + 1;
sz[u] = 1;
int max_sub = 0;
for (int v : adj[u]) {
if (v == p) continue;
dfs1(v, u);
sz[u] += sz[v];
if (sz[v] > max_sub) {
max_sub = sz[v];
heavy[u] = v;
}
}
};
dfs1(root, 0);
Second DFS (Chain Decomposition):
vector<int> top(n+1), pos(n+1);
int timer = 0;
function<void(int,int)> dfs2 = [&](int u, int tp) {
top[u] = tp;
pos[u] = ++timer;
if (heavy[u] != -1) dfs2(heavy[u], tp);
for (int v : adj[u])
if (v != parent[u] && v != heavy[u])
dfs2(v, v);
};
dfs2(root, root);
Path Operations Using Chains:
auto path_op = [&](int u, int v) {
while (top[u] != top[v]) {
if (depth[top[u]] < depth[top[v]]) swap(u, v);
// Operate on [pos[top[u]], pos[u]]
u = parent[top[u]];
}
if (depth[u] > depth[v]) swap(u, v);
// Operate on [pos[u], pos[v]]
};
LCA via HLD
auto lca = [&](int u, int v) {
while (top[u] != top[v]) {
if (depth[top[u]] < depth[top[v]]) swap(u, v);
u = parent[top[u]];
}
return depth[u] < depth[v] ? u : v;
};
Edge Weight Conversion
To handle edge weights, store each edge’s weight in its child node during the first DFS. Then, when querying the path between u and v, exclude the LCA by querying from pos[lca]+1 to pos[node].
Root Change Support
When the tree is rerooted, subtree operations require identifying the correct child of the original root that leads to the new root. This can be done by traversing upward until reaching the immediate child on the path to the new root.