CP-Library
This documentation is automatically generated by online-judge-tools/verification-helper
graph/tree/ReRooting.hpp
- View this file on GitHub
- Last update: 2026-08-25 17:03:15+07:00
- Include:
#include "graph/tree/ReRooting.hpp"
Depends on
graph/Graph.hpp
Code
#pragma once
#include <vector>
#include <utility>
#include "graph/Graph.hpp"
template <class M, class T, class F>
class ReRooting {
private:
using U = typename M::value_type;
const F& f;
int n;
const Graph<T>& G;
std::vector<U> init_data;
std::vector<std::vector<U>> dp;
std::vector<U> res;
std::vector<int> par;
void dfs1(int v, int p) {
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
if (e.to == p) {
par[v] = i;
} else {
dfs1(e.to, v);
}
}
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
if (e.to == p) continue;
dp[v][par[v]] = M::op(
dp[v][par[v]],
f(dp[e.to][par[e.to]], Edge<T>(e.to, v, e.cost, e.id))
);
}
if (p != -1 && G[v].size() == 1) {
dp[v][par[v]] = init_data[v];
}
}
void dfs2(int v, int p, int v_id) {
std::vector<U> memo(G[v].size());
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
memo[i] = f(dp[e.to][e.to == p ? v_id : par[e.to]],
Edge<T>(e.to, v, e.cost, e.id));
}
dp[v][G[v].size() - 1] = M::e();
for (int i = (int)G[v].size() - 2; i >= 0; --i) {
dp[v][i] = M::op(memo[i + 1], dp[v][i + 1]);
}
U sml = M::e();
for (int i = 0; i < (int)G[v].size(); ++i) {
dp[v][i] = M::op(sml, dp[v][i]);
sml = M::op(sml, std::move(memo[i]));
}
dp[v].back() = std::move(sml);
if (G[v].size() == 1) {
dp[v][p == -1 ? 0 : par[v]] = init_data[v];
}
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
if (e.to != p) {
dfs2(e.to, v, i);
}
}
}
void init() {
n = G.size();
if (n == 1) {
res = init_data;
dp.assign(1, std::vector<U>{});
return;
}
dp.resize(n);
for (int i = 0; i < n; ++i) {
dp[i].assign(G[i].size() + 1, M::e());
}
par.resize(n);
par[0] = G[0].size();
dfs1(0, -1);
dfs2(0, -1, -1);
res.resize(n);
for (int i = 0; i < n; ++i) {
res[i] = std::move(dp[i].back());
dp[i].pop_back();
}
}
public:
ReRooting(const Graph<T>& G, const F& f)
: f(f), n(G.size()), G(G), init_data(std::vector<U>(G.size(), M::e())) {
init();
}
ReRooting(const Graph<T>& G, const F& f, const std::vector<U>& ind)
: f(f), n(G.size()), G(G), init_data(ind) {
init();
}
ReRooting(const Graph<T>& G, const F& f, std::vector<U>&& ind)
: f(f), n(G.size()), G(G), init_data(std::move(ind)) {
init();
}
const std::vector<U>& get_res() const& { return res; }
std::vector<U> get_res() && { return std::move(res); }
const U& operator[](int v) const& { return res[v]; }
U operator[](int v) && { return std::move(res[v]); }
const std::vector<std::vector<U>>& get_dp() const& { return dp; }
std::vector<std::vector<U>> get_dp() && { return std::move(dp); }
const U& get_dp(int v, int p_idx) const& { return dp[v][p_idx]; }
U get_dp(int v, int p_idx) && { return std::move(dp[v][p_idx]); }
};
template <class M, class T, class F>
ReRooting<M, T, F> make_rerooting(const Graph<T>& G, const F& f) {
return ReRooting<M, T, F>(G, f);
}#line 2 "graph/tree/ReRooting.hpp"
#include <vector>
#include <utility>
#line 3 "graph/Graph.hpp"
#include <cassert>
#line 5 "graph/Graph.hpp"
template <typename T = int>
struct Edge {
int from, to;
T cost;
int id;
Edge() : from(-1), to(-1), cost(0), id(-1) {}
Edge(int _from, int _to, const T& _cost = 1, int _id = -1)
: from(_from), to(_to), cost(_cost), id(_id) {}
Edge(int _from, int _to, T&& _cost, int _id = -1)
: from(_from), to(_to), cost(std::move(_cost)), id(_id) {}
operator int() const { return to; }
friend bool operator<(const Edge<T>& lhs, const Edge<T>& rhs) {
return lhs.cost < rhs.cost;
}
friend bool operator>(const Edge<T>& lhs, const Edge<T>& rhs) {
return lhs.cost > rhs.cost;
}
};
template <typename T = int>
using Edges = std::vector<Edge<T>>;
template <typename T = int>
using GMatrix = std::vector<std::vector<T>>;
template <typename T = int>
struct Graph : std::vector<std::vector<Edge<T>>> {
using Base = std::vector<std::vector<Edge<T>>>;
using Base::Base;
int eid = 0;
int edge_size() const { return eid; }
int add_edge(int a, int b, const T& c = 1, bool is_directed = false) {
assert(0 <= a && a < (int)this->size());
assert(0 <= b && b < (int)this->size());
(*this)[a].emplace_back(a, b, c, eid);
if (!is_directed) {
(*this)[b].emplace_back(b, a, c, eid);
}
return eid++;
}
};
#line 5 "graph/tree/ReRooting.hpp"
template <class M, class T, class F>
class ReRooting {
private:
using U = typename M::value_type;
const F& f;
int n;
const Graph<T>& G;
std::vector<U> init_data;
std::vector<std::vector<U>> dp;
std::vector<U> res;
std::vector<int> par;
void dfs1(int v, int p) {
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
if (e.to == p) {
par[v] = i;
} else {
dfs1(e.to, v);
}
}
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
if (e.to == p) continue;
dp[v][par[v]] = M::op(
dp[v][par[v]],
f(dp[e.to][par[e.to]], Edge<T>(e.to, v, e.cost, e.id))
);
}
if (p != -1 && G[v].size() == 1) {
dp[v][par[v]] = init_data[v];
}
}
void dfs2(int v, int p, int v_id) {
std::vector<U> memo(G[v].size());
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
memo[i] = f(dp[e.to][e.to == p ? v_id : par[e.to]],
Edge<T>(e.to, v, e.cost, e.id));
}
dp[v][G[v].size() - 1] = M::e();
for (int i = (int)G[v].size() - 2; i >= 0; --i) {
dp[v][i] = M::op(memo[i + 1], dp[v][i + 1]);
}
U sml = M::e();
for (int i = 0; i < (int)G[v].size(); ++i) {
dp[v][i] = M::op(sml, dp[v][i]);
sml = M::op(sml, std::move(memo[i]));
}
dp[v].back() = std::move(sml);
if (G[v].size() == 1) {
dp[v][p == -1 ? 0 : par[v]] = init_data[v];
}
for (int i = 0; i < (int)G[v].size(); ++i) {
const auto& e = G[v][i];
if (e.to != p) {
dfs2(e.to, v, i);
}
}
}
void init() {
n = G.size();
if (n == 1) {
res = init_data;
dp.assign(1, std::vector<U>{});
return;
}
dp.resize(n);
for (int i = 0; i < n; ++i) {
dp[i].assign(G[i].size() + 1, M::e());
}
par.resize(n);
par[0] = G[0].size();
dfs1(0, -1);
dfs2(0, -1, -1);
res.resize(n);
for (int i = 0; i < n; ++i) {
res[i] = std::move(dp[i].back());
dp[i].pop_back();
}
}
public:
ReRooting(const Graph<T>& G, const F& f)
: f(f), n(G.size()), G(G), init_data(std::vector<U>(G.size(), M::e())) {
init();
}
ReRooting(const Graph<T>& G, const F& f, const std::vector<U>& ind)
: f(f), n(G.size()), G(G), init_data(ind) {
init();
}
ReRooting(const Graph<T>& G, const F& f, std::vector<U>&& ind)
: f(f), n(G.size()), G(G), init_data(std::move(ind)) {
init();
}
const std::vector<U>& get_res() const& { return res; }
std::vector<U> get_res() && { return std::move(res); }
const U& operator[](int v) const& { return res[v]; }
U operator[](int v) && { return std::move(res[v]); }
const std::vector<std::vector<U>>& get_dp() const& { return dp; }
std::vector<std::vector<U>> get_dp() && { return std::move(dp); }
const U& get_dp(int v, int p_idx) const& { return dp[v][p_idx]; }
U get_dp(int v, int p_idx) && { return std::move(dp[v][p_idx]); }
};
template <class M, class T, class F>
ReRooting<M, T, F> make_rerooting(const Graph<T>& G, const F& f) {
return ReRooting<M, T, F>(G, f);
}