グラフ/木/heavy-light decomposition
(test/graph/tree/heavy-light_decomposition.1.test.cpp)

View this file on GitHub
Last update: 2023-05-12 15:57:02+09:00
Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2667

Depends on

Code

/*
 * @title グラフ/木/heavy-light decomposition
 *
 * verification-helper: PROBLEM http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2667
 */

#include <functional>
#include <iostream>
#include <vector>

#include "emthrm/data_structure/fenwick_tree/fenwick_tree_supporting_range_add_query.hpp"
#include "emthrm/graph/edge.hpp"
#include "emthrm/graph/tree/heavy-light_decomposition.hpp"

int main() {
  int n, q;
  std::cin >> n >> q;
  std::vector<std::vector<emthrm::Edge<long long>>> graph(n);
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    std::cin >> a >> b;
    graph[a].emplace_back(a, b, 0);
    graph[b].emplace_back(b, a, 0);
  }
  emthrm::HeavyLightDecomposition<long long>
      heavy_light_decomposition(graph, 0);
  emthrm::FenwickTreeSupportingRangeAddQuery<long long> bit(n - 1);
  while (q--) {
    int type;
    std::cin >> type;
    if (type == 0) {
      int u, v;
      std::cin >> u >> v;
      std::cout << heavy_light_decomposition.query_e(
                       u, v,
                       [&bit](const int l, const int r) -> long long {
                         return bit.sum(l, r);
                       },
                       std::plus<long long>(), 0LL)
                << '\n';
    } else if (type == 1) {
      int v, x;
      std::cin >> v >> x;
      heavy_light_decomposition.update_subtree_e(
          v,
          [&bit, x](const int l, const int r) -> void {
            return bit.add(l, r, x);
          });
    }
  }
  return 0;
}

#line 1 "test/graph/tree/heavy-light_decomposition.1.test.cpp"
/*
 * @title グラフ/木/heavy-light decomposition
 *
 * verification-helper: PROBLEM http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2667
 */

#include <functional>
#include <iostream>
#include <vector>

#line 1 "include/emthrm/data_structure/fenwick_tree/fenwick_tree_supporting_range_add_query.hpp"



#line 5 "include/emthrm/data_structure/fenwick_tree/fenwick_tree_supporting_range_add_query.hpp"

namespace emthrm {

template <typename Abelian>
struct FenwickTreeSupportingRangeAddQuery {
  explicit FenwickTreeSupportingRangeAddQuery(
      const int n_, const Abelian ID = 0)
      : n(n_ + 1), ID(ID) {
    data_const.assign(n, ID);
    data_linear.assign(n, ID);
  }

  void add(int left, const int right, const Abelian val) {
    if (right < ++left) [[unlikely]] return;
    for (int i = left; i < n; i += i & -i) {
      data_const[i] -= val * (left - 1);
      data_linear[i] += val;
    }
    for (int i = right + 1; i < n; i += i & -i) {
      data_const[i] += val * right;
      data_linear[i] -= val;
    }
  }

  Abelian sum(const int idx) const {
    Abelian res = ID;
    for (int i = idx; i > 0; i -= i & -i) {
      res += data_linear[i];
    }
    res *= idx;
    for (int i = idx; i > 0; i -= i & -i) {
      res += data_const[i];
    }
    return res;
  }

  Abelian sum(const int left, const int right) const {
    return left < right ? sum(right) - sum(left) : ID;
  }

  Abelian operator[](const int idx) const { return sum(idx, idx + 1); }

 private:
  const int n;
  const Abelian ID;
  std::vector<Abelian> data_const, data_linear;
};

}  // namespace emthrm


#line 1 "include/emthrm/graph/edge.hpp"
/**
 * @title 辺
 */

#ifndef EMTHRM_GRAPH_EDGE_HPP_
#define EMTHRM_GRAPH_EDGE_HPP_

#include <compare>

namespace emthrm {

template <typename CostType>
struct Edge {
  CostType cost;
  int src, dst;

  explicit Edge(const int src, const int dst, const CostType cost = 0)
      : cost(cost), src(src), dst(dst) {}

  auto operator<=>(const Edge& x) const = default;
};

}  // namespace emthrm

#endif  // EMTHRM_GRAPH_EDGE_HPP_
#line 1 "include/emthrm/graph/tree/heavy-light_decomposition.hpp"



#include <algorithm>
#include <utility>
#line 7 "include/emthrm/graph/tree/heavy-light_decomposition.hpp"

#line 1 "include/emthrm/graph/edge.hpp"
/**
 * @title 辺
 */

#ifndef EMTHRM_GRAPH_EDGE_HPP_
#define EMTHRM_GRAPH_EDGE_HPP_

#include <compare>

namespace emthrm {

template <typename CostType>
struct Edge {
  CostType cost;
  int src, dst;

  explicit Edge(const int src, const int dst, const CostType cost = 0)
      : cost(cost), src(src), dst(dst) {}

  auto operator<=>(const Edge& x) const = default;
};

}  // namespace emthrm

#endif  // EMTHRM_GRAPH_EDGE_HPP_
#line 9 "include/emthrm/graph/tree/heavy-light_decomposition.hpp"

namespace emthrm {

template <typename CostType>
struct HeavyLightDecomposition {
  std::vector<int> parent, subtree, id, inv, head;
  std::vector<CostType> cost;

  explicit HeavyLightDecomposition(
      const std::vector<std::vector<Edge<CostType>>>& graph,
      const int root = 0)
      : graph(graph) {
    const int n = graph.size();
    parent.assign(n, -1);
    subtree.assign(n, 1);
    dfs1(root);
    id.resize(n);
    inv.resize(n);
    head.assign(n, root);
    int cur_id = 0;
    dfs2(root, &cur_id);
  }

  template <typename Fn>
  void update_v(int u, int v, const Fn f) const {
    while (true) {
      if (id[u] > id[v]) std::swap(u, v);
      f(std::max(id[head[v]], id[u]), id[v] + 1);
      if (head[u] == head[v]) break;
      v = parent[head[v]];
    }
  }

  template <typename F, typename G, typename T>
  T query_v(int u, int v, const F f, const G g, const T id_t) const {
    T left = id_t, right = id_t;
    while (true) {
      if (id_t[u] > id_t[v]) {
        std::swap(u, v);
        std::swap(left, right);
      }
      left = g(left, f(std::max(id_t[head[v]], id_t[u]), id_t[v] + 1));
      if (head[u] == head[v]) break;
      v = parent[head[v]];
    }
    return g(left, right);
  }

  template <typename Fn>
  void update_subtree_v(const int ver, const Fn f) const {
    f(id[ver], id[ver] + subtree[ver]);
  }

  template <typename T, typename Fn>
  T query_subtree_v(const int ver, const Fn f) const {
    return f(id[ver], id[ver] + subtree[ver]);
  }

  template <typename Fn>
  void update_e(int u, int v, const Fn f) const {
    while (true) {
      if (id[u] > id[v]) std::swap(u, v);
      if (head[u] == head[v]) {
        f(id[u], id[v]);
        break;
      } else {
        f(id[head[v]] - 1, id[v]);
        v = parent[head[v]];
      }
    }
  }

  template <typename F, typename G, typename T>
  T query_e(int u, int v, const F f, const G g, const T id_t) const {
    T left = id_t, right = id_t;
    while (true) {
      if (id[u] > id[v]) {
        std::swap(u, v);
        std::swap(left, right);
      }
      if (head[u] == head[v]) {
        left = g(left, f(id[u], id[v]));
        break;
      } else {
        left = g(left, f(id[head[v]] - 1, id[v]));
        v = parent[head[v]];
      }
    }
    return g(left, right);
  }

  template <typename Fn>
  void update_subtree_e(const int ver, const Fn f) const {
    f(id[ver], id[ver] + subtree[ver] - 1);
  }

  template <typename T, typename Fn>
  T query_subtree_e(const int ver, const Fn f) const {
    return f(id[ver], id[ver] + subtree[ver] - 1);
  }

  int lowest_common_ancestor(int u, int v) const {
    while (true) {
      if (id[u] > id[v]) std::swap(u, v);
      if (head[u] == head[v]) break;
      v = parent[head[v]];
    }
    return u;
  }

 private:
  std::vector<std::vector<Edge<CostType>>> graph;

  void dfs1(const int ver) {
    for (int i = 0; std::cmp_less(i, graph[ver].size()); ++i) {
      Edge<CostType>& e = graph[ver][i];
      if (e.dst != parent[ver]) {
        parent[e.dst] = ver;
        dfs1(e.dst);
        subtree[ver] += subtree[e.dst];
        if (subtree[e.dst] > subtree[graph[ver].front().dst]) {
          std::swap(e, graph[ver].front());
        }
      }
    }
  }

  void dfs2(const int ver, int* cur_id) {
    id[ver] = (*cur_id)++;
    inv[id[ver]] = ver;
    for (const Edge<CostType>& e : graph[ver]) {
      if (e.dst != parent[ver]) {
        head[e.dst] = (e.dst == graph[ver].front().dst ? head[ver] : e.dst);
        cost.emplace_back(e.cost);
        dfs2(e.dst, cur_id);
      }
    }
  }
};

}  // namespace emthrm


#line 14 "test/graph/tree/heavy-light_decomposition.1.test.cpp"

int main() {
  int n, q;
  std::cin >> n >> q;
  std::vector<std::vector<emthrm::Edge<long long>>> graph(n);
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    std::cin >> a >> b;
    graph[a].emplace_back(a, b, 0);
    graph[b].emplace_back(b, a, 0);
  }
  emthrm::HeavyLightDecomposition<long long>
      heavy_light_decomposition(graph, 0);
  emthrm::FenwickTreeSupportingRangeAddQuery<long long> bit(n - 1);
  while (q--) {
    int type;
    std::cin >> type;
    if (type == 0) {
      int u, v;
      std::cin >> u >> v;
      std::cout << heavy_light_decomposition.query_e(
                       u, v,
                       [&bit](const int l, const int r) -> long long {
                         return bit.sum(l, r);
                       },
                       std::plus<long long>(), 0LL)
                << '\n';
    } else if (type == 1) {
      int v, x;
      std::cin >> v >> x;
      heavy_light_decomposition.update_subtree_e(
          v,
          [&bit, x](const int l, const int r) -> void {
            return bit.add(l, r, x);
          });
    }
  }
  return 0;
}

グラフ/木/heavy-light decomposition (test/graph/tree/heavy-light_decomposition.1.test.cpp)

Depends on

Code

グラフ/木/heavy-light decomposition
(test/graph/tree/heavy-light_decomposition.1.test.cpp)