dsu_on_tree.hpp

SECTIONGraph INCLUDEnoya/dsu_on_tree.hpp

按重儿子保留统计状态，离线计算每个子树的答案；适合能逐点加入/删除贡献的子树询问。

Complexity: Time: O(n log n) callback invocations on a tree. Space: O(n).

跳到代码 · GitHub ↗

Implementation

ReferenceSourceExpanded

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(n log n) callback invocations on a tree.
/// Space: O(n).

#include "noya/heavy_light_decomposition.hpp"

namespace noya {

/// @brief DSU on tree. add(u) adds vertex u, query(u, x) queries contribution of x to u, del(u) removes vertex u.
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
                 auto &&del) {
  hld hl1(g);
  auto dfs = [&](auto &dfs, int u, bool so1) -> void {
    for (auto v : g[u]) {
      if (v != hl1.fa[u] && v != hl1.son[u]) {
        dfs(dfs, v, false);
      }
    }
    if (hl1.son[u] != -1)
      dfs(dfs, hl1.son[u], true);

    for (auto v : g[u]) {
      if (v != hl1.fa[u] && v != hl1.son[u]) {
        auto [l, r] = hl1.subtree(v);
        for (int i = l; i < r; i++) {
          int x = hl1.tl[i];
          qry(u, x);
        }
        for (int i = l; i < r; i++) {
          int x = hl1.tl[i];
          add(x);
        }
      }
    }

    qry(u, u);
    add(u);

    if (!so1) {
      auto [l, r] = hl1.subtree(u);
      for (int i = l; i < r; i++) {
        int x = hl1.tl[i];
        del(x);
      }
    }
  };
  dfs(dfs, 0, true);
}
} // namespace noya

#ifndef NOYA_DSU_ON_TREE_HPP
#define NOYA_DSU_ON_TREE_HPP 1

/// @complexity Time: O(n log n) callback invocations on a tree.
/// Space: O(n).

#include "noya/heavy_light_decomposition.hpp"

namespace noya {

/// @brief DSU on tree. add(u) adds vertex u, query(u, x) queries contribution of x to u, del(u) removes vertex u.
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
                 auto &&del) {
  hld hl1(g);
  auto dfs = [&](auto &dfs, int u, bool so1) -> void {
    for (auto v : g[u]) {
      if (v != hl1.fa[u] && v != hl1.son[u]) {
        dfs(dfs, v, false);
      }
    }
    if (hl1.son[u] != -1)
      dfs(dfs, hl1.son[u], true);

    for (auto v : g[u]) {
      if (v != hl1.fa[u] && v != hl1.son[u]) {
        auto [l, r] = hl1.subtree(v);
        for (int i = l; i < r; i++) {
          int x = hl1.tl[i];
          qry(u, x);
        }
        for (int i = l; i < r; i++) {
          int x = hl1.tl[i];
          add(x);
        }
      }
    }

    qry(u, u);
    add(u);

    if (!so1) {
      auto [l, r] = hl1.subtree(u);
      for (int i = l; i < r; i++) {
        int x = hl1.tl[i];
        del(x);
      }
    }
  };
  dfs(dfs, 0, true);
}
} // namespace noya

#endif // NOYA_DSU_ON_TREE_HPP

#include <algorithm>
#include <array>
#include <cassert>
#include <tuple>
#include <utility>
#include <vector>

/// @complexity Time: O(n log n) callback invocations on a tree.
/// Space: O(n).

/// @complexity Time: O(n) build, O(log n) LCA/path decomposition, O(1) subtree interval.
/// Space: O(n).

namespace noya {
/// @brief Heavy-light decomposition for path and subtree queries on trees.
struct hld {
  std::vector<std::vector<int>> G;
  int n;
  std::vector<int> dfn, siz, son, top, d, fa;
  int idx;
  void dfs(int u, int p) {
    siz[u] = 1;
    for (auto v : G[u])
      if (v != p) {
        fa[v] = u;
        d[v] = d[u] + 1;
        dfs(v, u), siz[u] += siz[v];
        if (son[u] == -1 || siz[v] > siz[son[u]])
          son[u] = v;
      }
  }

  std::vector<int> tl;
  void dfs2(int u, int t) {
    top[u] = t, dfn[u] = idx++;
    tl.push_back(u);
    if (son[u] != -1)
      dfs2(son[u], t);
    for (auto v : G[u])
      if (top[v] == -1)
        dfs2(v, v);
  }

  /// @brief Return the k-th ancestor of node a, or -1 if k > depth(a).
  int get_kth_ancestor(int a, int k) const {
    if (k < 0 || k >= d[a])
      return -1;

    int dst = d[a] - k;
    while (d[top[a]] > dst)
      a = fa[top[a]];

    int pos = dfn[a] - (d[a] - dst);
    return tl[pos];
  }

  /// @brief Return the k-th node (0-indexed) on the path from a to b, or -1 if out of range.
  int get_kth_node_on_path(int a, int b, int k) const {
    int anc = lca(a, b);
    int lhs = d[a] - d[anc];
    int rhs = d[b] - d[anc];

    if (k < 0 || k > lhs + rhs)
      return -1;

    if (k < lhs)
      return get_kth_ancestor(a, k);
    else
      return get_kth_ancestor(b, lhs + rhs - k);
  }

  hld(const std::vector<std::vector<int>> &g = {}, const int &rt = 0) {
    if (!g.empty())
      build(g, rt);
  }

  void build(const std::vector<std::vector<int>> &g = {}, const int &rt = 0) {
    n = g.size();
    G = g;
    siz.assign(n, 0);
    dfn.assign(n, -1);
    son.assign(n, -1);
    top.assign(n, -1);
    d.assign(n, 0);
    fa.assign(n, -1);
    d[rt] = 1;
    dfs(rt, -1);
    idx = 0;
    tl.clear();
    dfs2(rt, rt);
  }

  /// @brief Check if a is in the subtree of b.
  bool is_subtree(int a, int b) const {
    if (dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b]) {
      return true;
    } else {
      return false;
    }
  }

  /// @brief Return the lowest common ancestor of x and y.
  int lca(int x, int y) const {
    while (top[x] != top[y]) {
      if (d[top[x]] < d[top[y]])
        std::swap(x, y);
      x = fa[top[x]];
    }
    return d[x] < d[y] ? x : y;
  }

  /// @brief Decompose path x->y into chain segments. @return (dfn_l, dfn_r, direction).
  std::vector<std::tuple<int, int, bool>> chain(int x, int y) const {
    assert(0 <= x && x < n);
    assert(0 <= y && y < n);
    std::vector<std::tuple<int, int, bool>> L, R;
    while (top[x] != top[y]) {
      assert(0 <= x && x < n);
      assert(0 <= y && y < n);
      if (d[top[x]] > d[top[y]]) {
        L.emplace_back(dfn[top[x]], dfn[x] + 1, false);
        x = fa[top[x]];
      } else {
        R.emplace_back(dfn[top[y]], dfn[y] + 1, true);
        y = fa[top[y]];
      }
    }
    if (dfn[y] < dfn[x])
      L.emplace_back(dfn[y], dfn[x] + 1, false);
    else
      R.emplace_back(dfn[x], dfn[y] + 1, true);
    reverse(R.begin(), R.end());
    L.insert(L.end(), R.begin(), R.end());
    return L;
  }

  /// @brief Return the DFN range [l, r) for the subtree of node a.
  std::array<int, 2> subtree(int a) const { return {dfn[a], dfn[a] + siz[a]}; }

  /// @brief Return the LCA when the tree is re-rooted at c.
  int rooted_lca(int a, int b, int c) const {
    return lca(a, b) ^ lca(a, c) ^ lca(b, c);
  }

  /// @brief Compute the intersection of paths (a,b) and (c,d) as a pair of endpoints.
  std::pair<int, int> intersection(int a, int b, int c, int d) const {
    int ab = lca(a, b), ac = lca(a, c), ad = lca(a, d);
    int bc = lca(b, c), bd = lca(b, d), cd = lca(c, d);
    int x = ab ^ ac ^ bc;
    int y = ab ^ ad ^ bd;
    if (x != y) {
      return {x, y};
    }
    int z = ac ^ ad ^ cd;
    if (x != z) {
      x = -1;
    }
    return {x, x};
  }

  std::pair<int, int> intersection(std::pair<int, int> a,
                                   std::pair<int, int> b) const {
    return intersection(a.first, a.second, b.first, b.second);
  }
};
} // namespace noya

namespace noya {

/// @brief DSU on tree. add(u) adds vertex u, query(u, x) queries contribution of x to u, del(u) removes vertex u.
void dsu_on_tree(std::vector<std::vector<int>> &g, auto &&add, auto &&qry,
                 auto &&del) {
  hld hl1(g);
  auto dfs = [&](auto &dfs, int u, bool so1) -> void {
    for (auto v : g[u]) {
      if (v != hl1.fa[u] && v != hl1.son[u]) {
        dfs(dfs, v, false);
      }
    }
    if (hl1.son[u] != -1)
      dfs(dfs, hl1.son[u], true);

    for (auto v : g[u]) {
      if (v != hl1.fa[u] && v != hl1.son[u]) {
        auto [l, r] = hl1.subtree(v);
        for (int i = l; i < r; i++) {
          int x = hl1.tl[i];
          qry(u, x);
        }
        for (int i = l; i < r; i++) {
          int x = hl1.tl[i];
          add(x);
        }
      }
    }

    qry(u, u);
    add(u);

    if (!so1) {
      auto [l, r] = hl1.subtree(u);
      for (int i = l; i < r; i++) {
        int x = hl1.tl[i];
        del(x);
      }
    }
  };
  dfs(dfs, 0, true);
}
} // namespace noya