Merge pull request #364 from hitonanode/hld-subtree

HLD: subtree クエリに対応

Author: hitonanode

graph/extended_block_cut_trees.md

@@ -32,8 +32,8 @@ vector<vector<int>> bct_to = bct.to;
 assert(n + b == (int)bct_to.size());
 
 // 得られた block cut tree を heavy-light decomposition などに使う場合
-HeavyLightDecomposition hld(bct.size());
-for (auto [i, j] : bct.get_edges()) hld.add_edge(i, j);
+heavy_light_decomposition hld(bct.size(), bct.get_edges());
+hld.build();
 ```
 
 ## 問題例

graph/test/extended_block_cut_trees.yuki1326.test.cpp

@@ -19,12 +19,11 @@ int main() {
 
     const extended_block_cut_trees bct(N, edges);
 
-    HeavyLightDecomposition hld(bct.size());
-    for (auto [i, j] : bct.get_edges()) hld.add_edge(i, j);
+    heavy_light_decomposition hld(bct.size(), bct.get_edges());
     hld.build();
 
     atcoder::fenwick_tree<int> fw(hld.V);
-    for (int i = 0; i < N; ++i) fw.add(hld.aligned_id[i], 1);
+    for (int i = 0; i < N; ++i) fw.add(hld.subtree_begin[i], 1);
 
     int Q;
     cin >> Q;

tree/heavy_light_decomposition.hpp

@@ -7,99 +7,118 @@
 #include <utility>
 #include <vector>
 
-// Heavy-Light Decomposition of trees
+// Heavy-light decomposition of trees (forest)
 // Based on http://beet-aizu.hatenablog.com/entry/2017/12/12/235950
-struct HeavyLightDecomposition {
+struct heavy_light_decomposition {
     int V;
     int k;
-    int nb_heavy_path;
     std::vector<std::vector<int>> e;
     std::vector<int> par;         // par[i] = parent of vertex i (Default: -1)
     std::vector<int> depth;       // depth[i] = distance between root and vertex i
     std::vector<int> subtree_sz;  // subtree_sz[i] = size of subtree whose root is i
     std::vector<int> heavy_child; // heavy_child[i] = child of vertex i on heavy path (Default: -1)
     std::vector<int> tree_id;     // tree_id[i] = id of tree vertex i belongs to
-    std::vector<int> aligned_id,
-        aligned_id_inv;    // aligned_id[i] =  aligned id for vertex i (consecutive on heavy edges)
+    // subtree_begin[i] = aligned id for vertex i (consecutive on heavy edges)
+    std::vector<int> subtree_begin, subtree_end;
+    std::vector<int> vis_order; // vis_order[subtree_begin[i]] = i
     std::vector<int> head; // head[i] = id of vertex on heavy path of vertex i, nearest to root
-    std::vector<int> head_ids;      // consist of head vertex id's
-    std::vector<int> heavy_path_id; // heavy_path_id[i] = heavy_path_id for vertex [i]
-
-    HeavyLightDecomposition(int sz = 0)
-        : V(sz), k(0), nb_heavy_path(0), e(sz), par(sz), depth(sz), subtree_sz(sz), heavy_child(sz),
-          tree_id(sz, -1), aligned_id(sz), aligned_id_inv(sz), head(sz), heavy_path_id(sz, -1) {}
-    void add_edge(int u, int v) {
-        e[u].emplace_back(v);
-        e[v].emplace_back(u);
+
+    heavy_light_decomposition(int n, const std::vector<std::pair<int, int>> &edges)
+        : V(n), k(0), e(n), par(n), depth(n), subtree_sz(n), heavy_child(n), tree_id(n, -1),
+          subtree_begin(n), subtree_end(n), vis_order(0), head(n) {
+        for (auto [u, v] : edges) {
+            assert(u != v);
+            e.at(u).emplace_back(v);
+            e.at(v).emplace_back(u);
+        }
     }
 
-    void _build_dfs(int root) {
-        std::stack<std::pair<int, int>> st;
+    void _build_dfs_1(int root) {
+        std::vector<std::pair<int, int>> st;
         par[root] = -1;
         depth[root] = 0;
-        st.emplace(root, 0);
+        st.emplace_back(root, 0);
         while (!st.empty()) {
-            int now = st.top().first;
-            int &i = st.top().second;
+            int now = st.back().first;
+            int &i = st.back().second;
             if (i < (int)e[now].size()) {
                 int nxt = e[now][i++];
                 if (nxt == par[now]) continue;
                 par[nxt] = now;
                 depth[nxt] = depth[now] + 1;
-                st.emplace(nxt, 0);
+                st.emplace_back(nxt, 0);
             } else {
-                st.pop();
+                st.pop_back();
                 int max_sub_sz = 0;
                 subtree_sz[now] = 1;
                 heavy_child[now] = -1;
                 for (auto nxt : e[now]) {
                     if (nxt == par[now]) continue;
                     subtree_sz[now] += subtree_sz[nxt];
-                    if (max_sub_sz < subtree_sz[nxt])
+                    if (max_sub_sz < subtree_sz[nxt]) {
                         max_sub_sz = subtree_sz[nxt], heavy_child[now] = nxt;
+                    }
                 }
             }
         }
     }
 
-    void _build_bfs(int root, int tree_id_now) {
-        std::queue<int> q({root});
-        while (!q.empty()) {
-            int h = q.front();
-            q.pop();
-            head_ids.emplace_back(h);
-            for (int now = h; now != -1; now = heavy_child[now]) {
+    void _build_dfs_2(int now, int tree_id_now) {
+        std::vector<std::pair<int, bool>> st;
+        st.emplace_back(now, true);
+        head[now] = now;
+
+        while (!st.empty()) {
+            auto [now, mode] = st.back();
+            st.pop_back();
+
+            if (mode) {
                 tree_id[now] = tree_id_now;
-                aligned_id[now] = k++;
-                aligned_id_inv[aligned_id[now]] = now;
-                heavy_path_id[now] = nb_heavy_path;
-                head[now] = h;
-                for (int nxt : e[now])
-                    if (nxt != par[now] and nxt != heavy_child[now]) q.push(nxt);
+                subtree_begin[now] = k++;
+                vis_order.push_back(now);
+
+                st.emplace_back(now, false);
+
+                for (int nxt : e[now]) {
+                    if (nxt == par[now] or nxt == heavy_child[now]) continue;
+                    head[nxt] = nxt;
+                    st.emplace_back(nxt, true);
+                }
+
+                if (heavy_child[now] != -1) {
+                    head[heavy_child[now]] = head[now];
+                    st.emplace_back(heavy_child[now], true);
+                }
+            } else {
+                subtree_end[now] = k;
             }
-            nb_heavy_path++;
         }
     }
 
     void build(std::vector<int> roots = {0}) {
         int tree_id_now = 0;
-        for (auto r : roots) _build_dfs(r), _build_bfs(r, tree_id_now++);
+        for (int r : roots) {
+            _build_dfs_1(r);
+            _build_dfs_2(r, tree_id_now++);
+        }
     }
 
     template <class T> std::vector<T> segtree_rearrange(const std::vector<T> &data) const {
         assert(int(data.size()) == V);
         std::vector<T> ret;
-        ret.reserve(V);
-        for (int i = 0; i < V; i++) ret.emplace_back(data[aligned_id_inv[i]]);
+        ret.reserve(vis_order.size());
+        for (int v : vis_order) ret.emplace_back(data[v]);
         return ret;
     }
 
+    void for_vertex(int u, const std::function<void(int)> &f) const { f(subtree_begin[u]); }
+
     // query for vertices on path [u, v] (INCLUSIVE)
     void
     for_each_vertex(int u, int v, const std::function<void(int ancestor, int descendant)> &f) const {
         while (true) {
-            if (aligned_id[u] > aligned_id[v]) std::swap(u, v);
-            f(std::max(aligned_id[head[v]], aligned_id[u]), aligned_id[v]);
+            if (subtree_begin[u] > subtree_begin[v]) std::swap(u, v);
+            f(std::max(subtree_begin[head[v]], subtree_begin[u]), subtree_begin[v]);
             if (head[u] == head[v]) break;
             v = par[head[v]];
         }
@@ -112,7 +131,7 @@ struct HeavyLightDecomposition {
         const int lca = lowest_common_ancestor(u, v), dlca = depth[lca];
         while (u >= 0 and depth[u] > dlca) {
             const int p = (depth[head[u]] > dlca ? head[u] : lca);
-            fup(aligned_id[p] + (p == lca), aligned_id[u]), u = par[p];
+            fup(subtree_begin[p] + (p == lca), subtree_begin[u]), u = par[p];
         }
         static std::vector<std::pair<int, int>> lrs;
         int sz = 0;
@@ -121,28 +140,33 @@ struct HeavyLightDecomposition {
             if (int(lrs.size()) == sz) lrs.emplace_back(0, 0);
             lrs.at(sz++) = {p, v}, v = par.at(p);
         }
-        while (sz--) fdown(aligned_id[lrs.at(sz).first], aligned_id[lrs.at(sz).second]);
+        while (sz--) fdown(subtree_begin[lrs.at(sz).first], subtree_begin[lrs.at(sz).second]);
     }
 
     // query for edges on path [u, v]
     void for_each_edge(int u, int v, const std::function<void(int, int)> &f) const {
         while (true) {
-            if (aligned_id[u] > aligned_id[v]) std::swap(u, v);
+            if (subtree_begin[u] > subtree_begin[v]) std::swap(u, v);
             if (head[u] != head[v]) {
-                f(aligned_id[head[v]], aligned_id[v]);
+                f(subtree_begin[head[v]], subtree_begin[v]);
                 v = par[head[v]];
             } else {
-                if (u != v) f(aligned_id[u] + 1, aligned_id[v]);
+                if (u != v) f(subtree_begin[u] + 1, subtree_begin[v]);
                 break;
             }
         }
     }
 
+    // query for vertices in subtree rooted at u
+    void for_subtree(int u, const std::function<void(int, int)> &f) const {
+        f(subtree_begin[u], subtree_end[u] - 1);
+    }
+
     // lowest_common_ancestor: O(log V)
     int lowest_common_ancestor(int u, int v) const {
         assert(tree_id[u] == tree_id[v] and tree_id[u] >= 0);
         while (true) {
-            if (aligned_id[u] > aligned_id[v]) std::swap(u, v);
+            if (subtree_begin[u] > subtree_begin[v]) std::swap(u, v);
             if (head[u] == head[v]) return u;
             v = par[head[v]];
         }
@@ -159,7 +183,7 @@ struct HeavyLightDecomposition {
         if (k < 0) return -1;
         while (v >= 0) {
             int h = head.at(v), len = depth.at(v) - depth.at(h);
-            if (k <= len) return aligned_id_inv.at(aligned_id.at(v) - k);
+            if (k <= len) return vis_order.at(subtree_begin.at(v) - k);
             k -= len + 1, v = par.at(h);
         }
         return -1;

tree/heavy_light_decomposition.md

@@ -3,39 +3,36 @@ title: Heavy-light decomposition （HLD, 木の重軽分解）
 documentation_of: ./heavy_light_decomposition.hpp
 ---
 
-木の重軽分解を行い，列に対する処理を木上のパスに拡張した処理や，最小共通祖先の効率的な計算を可能にする．前処理 $O(N)$．
+木の重軽分解を行い，列に対する処理を木上のパスや部分木に含まれる全ての頂点に拡張した処理や，最小共通祖先の効率的な計算を可能にする．前処理 $O(n)$．
 
 ## 使用方法
 
 ### 前処理
 
 ```cpp
-HeavyLightDecomposition hld(N);
-for (int e = 0; e < N - 1; ++e) {
-    int u, v;
-    cin >> u >> v;
-    --u, --v;
-    hld.add_edge(u, v);
-}
+int n;  // 頂点数
+vector<pair<int, int>> edges;
+
+heavy_light_decomposition hld(n, edges);
 hld.build();  // O(N)
 ```
 
 ### 木に対する基本的なクエリ
 
-以下，すべてクエリ $O(\log N)$．
+以下，すべてクエリ $O(\log n)$．
 
 ```cpp
 cout << hld.lowest_common_ancestor(a, b) << '\n'; // (a, b) の最長共通祖先
 cout << hld.kth_parent(a, k) << '\n';  // 頂点 a の k 世代親
-cout << hld.distance(a, b) << '\n'; // 2 頂点 a, b の距離
+cout << hld.distance(a, b) << '\n';  // 2 頂点 a, b の距離
 cout << hld.s_to_t_by_k_steps(s, t, k) << '\n'; // 頂点 s から t 方向に k 歩動いた頂点
 ```
 
 ### セグメント木などを利用した木上パスクエリへの対応
 
 ```cpp
 // 各頂点の情報
-std::vector<S> vertices_info(N);
+std::vector<S> vertices_info(n);
 
 // HLD の分解順を考慮したセグメント木の初期化
 // セグメント木上で，分解されたパス上の要素は「根方向のものから順に」連続して並ぶことに注意
@@ -50,3 +47,10 @@ hld.for_each_vertex(from, to, func);
 ## 問題例
 
 - [Library Checker: Lowest Common Ancestor](https://judge.yosupo.jp/problem/lca)
+- [Library Checker: Vertex Add Subtree Sum](https://judge.yosupo.jp/problem/vertex_add_subtree_sum)
+- [いろはちゃんコンテスト Day2 K - 虫取り](https://atcoder.jp/contests/iroha2019-day2/tasks/iroha2019_day2_k)
+  - 部分木クエリとパスクエリの両方が使われる．
+
+## 参考文献・リンク
+
+- [Easiest HLD with subtree queries - Codeforces](https://codeforces.com/blog/entry/53170)

tree/test/hl_decomposition.test.cpp

@@ -1,19 +1,25 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/lca"
 #include "../heavy_light_decomposition.hpp"
 #include <iostream>
-#define PROBLEM "https://judge.yosupo.jp/problem/lca"
+using namespace std;
 
 int main() {
+    cin.tie(nullptr);
+    ios::sync_with_stdio(false);
+
     int N, Q, p, u, v;
-    std::cin >> N >> Q;
-    HeavyLightDecomposition hld(N);
+    cin >> N >> Q;
+    vector<pair<int, int>> edges;
     for (int i = 1; i <= N - 1; i++) {
-        std::cin >> p;
-        hld.add_edge(i, p);
+        cin >> p;
+        edges.emplace_back(i, p);
     }
+
+    heavy_light_decomposition hld(N, edges);
     hld.build();
 
     while (Q--) {
-        std::cin >> u >> v;
-        std::cout << hld.lowest_common_ancestor(u, v) << "\n";
+        cin >> u >> v;
+        cout << hld.lowest_common_ancestor(u, v) << "\n";
     }
 }

tree/test/hld_subtree_query.test.cpp

@@ -0,0 +1,46 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"
+#include "../heavy_light_decomposition.hpp"
+#include <iostream>
+#include <utility>
+#include <vector>
+using namespace std;
+
+#include <atcoder/fenwicktree>
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+
+    int N, Q;
+    cin >> N >> Q;
+    vector<long long> A(N);
+    for (auto &x : A) cin >> x;
+
+    vector<pair<int, int>> edges;
+    for (int i = 1; i < N; ++i) {
+        int p;
+        cin >> p;
+        edges.emplace_back(i, p);
+    }
+
+    heavy_light_decomposition hld(N, edges);
+    hld.build();
+
+    atcoder::fenwick_tree<long long> ft(N);
+
+    A = hld.segtree_rearrange(A);
+    for (int i = 0; i < N; ++i) ft.add(i, A.at(i));
+
+    while (Q--) {
+        int tp, u;
+        cin >> tp >> u;
+        if (tp == 0) {
+            long long x;
+            cin >> x;
+            hld.for_vertex(u, [&](int i) { ft.add(i, x); });
+        } else {
+            long long ans = 0;
+            hld.for_subtree(u, [&](int l, int r) { ans += ft.sum(l, r + 1); });
+            cout << ans << '\n';
+        }
+    }
+}

tree/test/jump_on_tree_hld.test.cpp

@@ -7,13 +7,14 @@ int main() {
     cin.tie(nullptr), ios::sync_with_stdio(false);
     int N, Q;
     cin >> N >> Q;
-    HeavyLightDecomposition tree(N);
+    vector<pair<int, int>> edges;
     for (int e = 0; e < N - 1; ++e) {
         int a, b;
         cin >> a >> b;
-        tree.add_edge(a, b);
+        edges.emplace_back(a, b);
     }
 
+    heavy_light_decomposition tree(N, edges);
     tree.build({0});
 
     while (Q--) {