cp-library
競技プログラミング用 C++ ライブラリ
全方位木 DP
tree/rerooting.hpp · GitHub で見る
共通ページ
全頂点をそれぞれ根にした木 DP を計算します。辺の両方向の寄与も取得できます。次は、各頂点から全頂点までの距離の総和を求める例です。
#include "lib/tree/rerooting.hpp"
using Data = std::pair<long long, int>;
Rerooting<Data> tree(3);
std::vector<long long> weight{2, 5};
tree.add_edge(0, 1);
tree.add_edge(1, 2);
const auto &answer = tree.build(Data{0, 0},
[](Data a, Data b) {
return Data{a.first + b.first, a.second + b.second};
},
[](Data x, int) {
return Data{x.first, x.second + 1};
},
[&](Data x, int, int, int id) {
return Data{x.first + weight[id] * x.second, x.second};
});
long long sum = answer[1].first; // 7
auto from_zero = tree.edge_dp(0); // {2, 1}
T は頂点の DP 値、E は辺を通した寄与の型です。E を省略すると T と同じ型になります。Data::first が距離和、second が頂点数です。辺重みや頂点の値は、コールバックから参照します。
| 関数・型 | 内容・返り値 | 計算量 |
|---|---|---|
Rerooting<T,E = T>(n = 0) |
n 頂点の空の森を作る |
\(O(n)\) |
add_edge(u,v) |
無向辺を追加し、0 始まりの辺 ID を返す | 償却 \(O(1)\) |
build(unit,merge,add_vertex,lift,root = 0) |
全方位木 DP を計算し、全頂点の答えを const vector<T>& で返す |
\(O(n)\) |
tree[v] |
頂点 v を根とした答えを const T& で返す |
\(O(1)\) |
values() |
全頂点の答えを const vector<T>& で返す |
\(O(1)\) |
edge_dp(id,reverse = false) |
辺を通した寄与を const E& で返す。add_edge(u,v) に対して false は u → v、true は v → u
|
\(O(1)\) |
edge(id) |
入力した辺の両端 (u,v) を const pair<int,int>& で返す |
\(O(1)\) |
size() / empty() / edge_count()
|
頂点数／頂点数が 0 か／辺数 | \(O(1)\) |
Vertex / Edge
|
それぞれ T / E の型別名 |
\(O(1)\) |
| コールバック | 内容・返り値 |
|---|---|
merge(const E& a,const E& b) -> E |
隣接する成分の寄与を結合する。結合則と両側の単位元 unit が必要。可換性・逆元は不要 |
add_vertex(const E& x,int v) -> T |
頂点 v に集まった寄与に v 自身を加える |
lift(const T& x,int from,int to,int id) -> E |
頂点 from を根とした値を、辺 id を通して to 側の寄与に変換する |
edge_dp(id) は、その辺を削除したときの送信元側の成分に対し、送信元の頂点を加え、最後に辺を通す変換を施した値です。受信先の頂点は含みません。例の from_zero は頂点 0 だけを含み、頂点 1 から測った距離が 2、頂点数が 1 です。
同じ頂点の隣接辺は add_edge の呼び出し順に結合します。親への辺を除くときも、残った辺の順序を変えません。各コールバックは引数と入力データだけで返り値が決まるようにします。root は内部の走査を始める頂点で、答えには影響しません。再帰は使いません。
入力は 0 始まりの木・森とし、自己ループ・多重辺・閉路を含めないでください。森では各頂点の属する連結成分について計算します。孤立点の答えは add_vertex(unit,v)、空の森の答えは空配列です。空でない場合は \(0\le\texttt{root}<n\) とします。
頂点数は INT_MAX 未満、辺数は INT_MAX/2 以下とします。
値型はコピー・代入可能としてください。T と E が異なる場合は T のデフォルト構築も必要です。戻り値をコピーしたくない場合は const auto& で受け取ります。tree[v]・values()・edge_dp は build 後に使います。辺の追加や再構築後は、以前取得した参照を使わないでください。入力データやコールバックを変えて build を呼び直すこともできます。
頂点数を \(n\) とし、各コールバック・値のコピーを \(O(1)\) とすると、構築・build は \(O(n)\) 時間・空間です。add_edge は償却 \(O(1)\)、そのほかの取得は \(O(1)\) です。
依存ライブラリ
コード
#ifndef REROOTING_HPP
#define REROOTING_HPP
#include "../components/undirected_csr.hpp"
#include <algorithm>
#include <cassert>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>
template<class T, class E = T>
class Rerooting {
int n_;
bool built_ = false;
std::vector<std::pair<int, int>> edges_;
std::vector<T> values_;
std::vector<E> messages_;
public:
using Vertex = T;
using Edge = E;
explicit Rerooting(int n = 0) : n_(n) {
assert(0 <= n && n < std::numeric_limits<int>::max());
edges_.reserve(std::max(0, n - 1));
}
int size() const { return n_; }
bool empty() const { return n_ == 0; }
int edge_count() const { return static_cast<int>(edges_.size()); }
int add_edge(int u, int v) {
assert(0 <= u && u < n_ && 0 <= v && v < n_ && u != v);
assert(edges_.size() < std::size_t(std::numeric_limits<int>::max() / 2));
built_ = false;
edges_.emplace_back(u, v);
return edge_count() - 1;
}
const std::pair<int, int>& edge(int id) const {
assert(0 <= id && id < edge_count());
return edges_[id];
}
const T& operator[](int v) const {
assert(built_ && 0 <= v && v < n_);
return values_[v];
}
const std::vector<T>& values() const {
assert(built_);
return values_;
}
const E& edge_dp(int id, bool reverse = false) const {
assert(built_ && 0 <= id && id < edge_count());
return messages_[2 * id + reverse];
}
template<class Merge, class AddVertex, class AddEdge>
const std::vector<T>& build(E unit, Merge&& merge, AddVertex&& add_vertex, AddEdge&& add_edge, int root = 0) {
assert(n_ == 0 || (0 <= root && root < n_));
built_ = false;
if constexpr (std::is_default_constructible_v<E>) messages_.resize(2 * edges_.size());
else messages_.assign(2 * edges_.size(), unit);
if (n_ == 0) {
values_.clear();
built_ = true;
return values_;
}
UndirectedCSR graph(n_, edges_);
for (int v = 0; v < n_; ++v)
for (int i = graph.start[v]; i < graph.start[v + 1]; ++i) {
int id = graph.adj[i].id;
graph.adj[i].id = 2 * id + (edges_[id].second == v);
}
std::vector<int> parent(n_, -2), parent_edge(n_, -1), order;
order.reserve(n_);
auto visit = [&](int s) {
if (parent[s] != -2) return;
parent[s] = -1;
order.push_back(s);
for (std::size_t i = order.size() - 1; i < order.size(); ++i) {
int v = order[i];
for (auto [to, id] : graph[v]) {
if (id == parent_edge[v]) continue;
assert(parent[to] == -2);
if (parent[to] != -2) continue;
parent[to] = v, parent_edge[to] = id ^ 1;
order.push_back(to);
}
}
};
visit(root);
for (int v = 0; v < n_; ++v) visit(v);
for (auto it = order.rbegin(); it != order.rend(); ++it) {
int v = *it, id = parent_edge[v];
if (id == -1) continue;
E value = unit;
for (auto arc : graph[v])
if (arc.id != id) value = merge(value, messages_[arc.id ^ 1]);
messages_[id] = add_edge(add_vertex(value, v), v, parent[v], id >> 1);
}
if constexpr (std::is_default_constructible_v<T>) values_.resize(n_);
else values_.assign(n_, unit);
int maximum = 0;
for (int v = 0; v < n_; ++v) maximum = std::max(maximum, graph.start[v + 1] - graph.start[v]);
std::vector<E> prefix(maximum, unit);
for (int v : order) {
auto arcs = graph[v];
int degree = static_cast<int>(arcs.size());
if (degree == 0) {
values_[v] = add_vertex(unit, v);
continue;
}
auto incoming = [&](int i) -> const E& { return messages_[arcs[i].id ^ 1]; };
if (degree > 1) prefix[1] = incoming(0);
for (int i = 1; i + 1 < degree; ++i) prefix[i + 1] = merge(prefix[i], incoming(i));
values_[v] = add_vertex(degree == 1 ? incoming(0) : merge(prefix[degree - 1], incoming(degree - 1)), v);
E suffix = unit;
for (int i = degree - 1; i >= 0; --i) {
auto [to, id] = arcs[i];
if (id != parent_edge[v]) {
E value = i == degree - 1 ? prefix[i] : i == 0 ? suffix : merge(prefix[i], suffix);
messages_[id] = add_edge(add_vertex(value, v), v, to, id >> 1);
}
if (i > 0) suffix = i == degree - 1 ? incoming(i) : merge(incoming(i), suffix);
}
}
built_ = true;
return values_;
}
};
#endif
#line 1 "tree/rerooting.hpp"
#line 1 "components/undirected_csr.hpp"
#include <cassert>
#include <cstddef>
#include <limits>
#include <memory>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
namespace undirected_csr_detail {
template <class Arc>
void build(int n, std::span<const std::pair<int, int>> edges, std::vector<int>& start, std::unique_ptr<Arc[]>& adj) {
assert(0 <= n && n < std::numeric_limits<int>::max());
assert(edges.size() <= std::size_t(std::numeric_limits<int>::max() / 2));
start.assign(std::size_t(n) + 1, 0);
for (auto [u, v] : edges) {
assert(0 <= u && u < n && 0 <= v && v < n);
++start[u + 1], ++start[v + 1];
}
for (int v = 0; v < n; ++v) start[v + 1] += start[v];
adj.reset(new Arc[2 * edges.size()]);
std::vector<int> cursor(start.begin(), start.end() - 1);
for (int i = 0; i < static_cast<int>(edges.size()); ++i) {
auto [u, v] = edges[i];
if constexpr (std::is_same_v<Arc, int>) adj[cursor[u]++] = v, adj[cursor[v]++] = u;
else adj[cursor[u]++] = {v, i}, adj[cursor[v]++] = {u, i};
}
}
}
struct UndirectedCSR {
struct Arc { int to, id; };
int n;
std::vector<int> start;
std::unique_ptr<Arc[]> adj;
UndirectedCSR(int n, std::span<const std::pair<int, int>> edges) : n(n) {
undirected_csr_detail::build(n, edges, start, adj);
}
std::span<const Arc> operator[](int v) const {
assert(0 <= v && v < n);
return {adj.get() + start[v], std::size_t(start[v + 1] - start[v])};
}
};
struct UndirectedVertexCSR {
int n;
std::vector<int> start;
std::unique_ptr<int[]> adj;
UndirectedVertexCSR(int n, std::span<const std::pair<int, int>> edges) : n(n) {
undirected_csr_detail::build(n, edges, start, adj);
}
std::span<const int> operator[](int v) const {
assert(0 <= v && v < n);
return {adj.get() + start[v], std::size_t(start[v + 1] - start[v])};
}
};
#line 5 "tree/rerooting.hpp"
#include <algorithm>
#line 11 "tree/rerooting.hpp"
template<class T, class E = T>
class Rerooting {
int n_;
bool built_ = false;
std::vector<std::pair<int, int>> edges_;
std::vector<T> values_;
std::vector<E> messages_;
public:
using Vertex = T;
using Edge = E;
explicit Rerooting(int n = 0) : n_(n) {
assert(0 <= n && n < std::numeric_limits<int>::max());
edges_.reserve(std::max(0, n - 1));
}
int size() const { return n_; }
bool empty() const { return n_ == 0; }
int edge_count() const { return static_cast<int>(edges_.size()); }
int add_edge(int u, int v) {
assert(0 <= u && u < n_ && 0 <= v && v < n_ && u != v);
assert(edges_.size() < std::size_t(std::numeric_limits<int>::max() / 2));
built_ = false;
edges_.emplace_back(u, v);
return edge_count() - 1;
}
const std::pair<int, int>& edge(int id) const {
assert(0 <= id && id < edge_count());
return edges_[id];
}
const T& operator[](int v) const {
assert(built_ && 0 <= v && v < n_);
return values_[v];
}
const std::vector<T>& values() const {
assert(built_);
return values_;
}
const E& edge_dp(int id, bool reverse = false) const {
assert(built_ && 0 <= id && id < edge_count());
return messages_[2 * id + reverse];
}
template<class Merge, class AddVertex, class AddEdge>
const std::vector<T>& build(E unit, Merge&& merge, AddVertex&& add_vertex, AddEdge&& add_edge, int root = 0) {
assert(n_ == 0 || (0 <= root && root < n_));
built_ = false;
if constexpr (std::is_default_constructible_v<E>) messages_.resize(2 * edges_.size());
else messages_.assign(2 * edges_.size(), unit);
if (n_ == 0) {
values_.clear();
built_ = true;
return values_;
}
UndirectedCSR graph(n_, edges_);
for (int v = 0; v < n_; ++v)
for (int i = graph.start[v]; i < graph.start[v + 1]; ++i) {
int id = graph.adj[i].id;
graph.adj[i].id = 2 * id + (edges_[id].second == v);
}
std::vector<int> parent(n_, -2), parent_edge(n_, -1), order;
order.reserve(n_);
auto visit = [&](int s) {
if (parent[s] != -2) return;
parent[s] = -1;
order.push_back(s);
for (std::size_t i = order.size() - 1; i < order.size(); ++i) {
int v = order[i];
for (auto [to, id] : graph[v]) {
if (id == parent_edge[v]) continue;
assert(parent[to] == -2);
if (parent[to] != -2) continue;
parent[to] = v, parent_edge[to] = id ^ 1;
order.push_back(to);
}
}
};
visit(root);
for (int v = 0; v < n_; ++v) visit(v);
for (auto it = order.rbegin(); it != order.rend(); ++it) {
int v = *it, id = parent_edge[v];
if (id == -1) continue;
E value = unit;
for (auto arc : graph[v])
if (arc.id != id) value = merge(value, messages_[arc.id ^ 1]);
messages_[id] = add_edge(add_vertex(value, v), v, parent[v], id >> 1);
}
if constexpr (std::is_default_constructible_v<T>) values_.resize(n_);
else values_.assign(n_, unit);
int maximum = 0;
for (int v = 0; v < n_; ++v) maximum = std::max(maximum, graph.start[v + 1] - graph.start[v]);
std::vector<E> prefix(maximum, unit);
for (int v : order) {
auto arcs = graph[v];
int degree = static_cast<int>(arcs.size());
if (degree == 0) {
values_[v] = add_vertex(unit, v);
continue;
}
auto incoming = [&](int i) -> const E& { return messages_[arcs[i].id ^ 1]; };
if (degree > 1) prefix[1] = incoming(0);
for (int i = 1; i + 1 < degree; ++i) prefix[i + 1] = merge(prefix[i], incoming(i));
values_[v] = add_vertex(degree == 1 ? incoming(0) : merge(prefix[degree - 1], incoming(degree - 1)), v);
E suffix = unit;
for (int i = degree - 1; i >= 0; --i) {
auto [to, id] = arcs[i];
if (id != parent_edge[v]) {
E value = i == degree - 1 ? prefix[i] : i == 0 ? suffix : merge(prefix[i], suffix);
messages_[id] = add_edge(add_vertex(value, v), v, to, id >> 1);
}
if (i > 0) suffix = i == degree - 1 ? incoming(i) : merge(incoming(i), suffix);
}
}
built_ = true;
return values_;
}
};