#pragma once
#include <vector>
#include "Mylib/Graph/Template/graph.cpp"
namespace haar_lib {
namespace rerooting_impl {
template <typename T, typename U, typename Merge, typename EdgeF, typename VertexF>
T rec1(
tree<U> &tr,
T id,
const Merge &merge,
const EdgeF &f,
const VertexF &g,
std::vector<std::vector<T>> &dp,
int cur,
int par = -1) {
T acc = id;
for (int i = 0; i < (int) tr[cur].size(); ++i) {
auto &e = tr[cur][i];
if (e.to == par) continue;
dp[cur][i] = rec1(tr, id, merge, f, g, dp, e.to, cur);
acc = merge(acc, f(dp[cur][i], e));
}
return g(acc, cur);
}
template <typename T, typename U, typename Merge, typename EdgeF, typename VertexF>
void rec2(
const tree<U> &tr,
T id,
const Merge &merge,
const EdgeF &f,
const VertexF &g,
std::vector<std::vector<T>> &dp,
int cur,
int par,
T value) {
const int l = tr[cur].size();
for (int i = 0; i < l; ++i) {
if (tr[cur][i].to == par) {
dp[cur][i] = value;
}
}
std::vector<T> left(l + 1, id), right(l + 1, id);
for (int i = 0; i < l - 1; ++i) {
const auto &e = tr[cur][i];
left[i + 1] = merge(left[i], f(dp[cur][i], e));
}
for (int i = l - 1; i >= 1; --i) {
const auto &e = tr[cur][i];
right[i - 1] = merge(right[i], f(dp[cur][i], e));
}
for (int i = 0; i < l; ++i) {
const auto &e = tr[cur][i];
if (e.to == par) continue;
rec2(tr, id, merge, f, g, dp, e.to, cur, g(merge(left[i], right[i]), cur));
}
}
} // namespace rerooting_impl
template <typename T, typename U, typename Merge, typename EdgeF, typename VertexF>
auto rerooting(tree<U> tr, T id, Merge merge, EdgeF f, VertexF g) {
const int N = tr.size();
std::vector<std::vector<T>> dp(N);
std::vector<T> ret(N, id);
for (int i = 0; i < N; ++i) dp[i].assign(tr[i].size(), id);
rerooting_impl::rec1(tr, id, merge, f, g, dp, 0);
rerooting_impl::rec2(tr, id, merge, f, g, dp, 0, -1, id);
for (int i = 0; i < N; ++i) {
for (int j = 0; j < (int) tr[i].size(); ++j) {
ret[i] = merge(ret[i], f(dp[i][j], tr[i][j]));
}
ret[i] = g(ret[i], i);
}
return ret;
}
} // namespace haar_lib
#line 2 "Mylib/Graph/TreeUtils/rerooting.cpp"
#include <vector>
#line 2 "Mylib/Graph/Template/graph.cpp"
#include <iostream>
#line 4 "Mylib/Graph/Template/graph.cpp"
namespace haar_lib {
template <typename T>
struct edge {
int from, to;
T cost;
int index = -1;
edge() {}
edge(int from, int to, T cost) : from(from), to(to), cost(cost) {}
edge(int from, int to, T cost, int index) : from(from), to(to), cost(cost), index(index) {}
};
template <typename T>
struct graph {
using weight_type = T;
using edge_type = edge<T>;
std::vector<std::vector<edge<T>>> data;
auto& operator[](size_t i) { return data[i]; }
const auto& operator[](size_t i) const { return data[i]; }
auto begin() const { return data.begin(); }
auto end() const { return data.end(); }
graph() {}
graph(int N) : data(N) {}
bool empty() const { return data.empty(); }
int size() const { return data.size(); }
void add_edge(int i, int j, T w, int index = -1) {
data[i].emplace_back(i, j, w, index);
}
void add_undirected(int i, int j, T w, int index = -1) {
add_edge(i, j, w, index);
add_edge(j, i, w, index);
}
template <size_t I, bool DIRECTED = true, bool WEIGHTED = true>
void read(int M) {
for (int i = 0; i < M; ++i) {
int u, v;
std::cin >> u >> v;
u -= I;
v -= I;
T w = 1;
if (WEIGHTED) std::cin >> w;
if (DIRECTED)
add_edge(u, v, w, i);
else
add_undirected(u, v, w, i);
}
}
};
template <typename T>
using tree = graph<T>;
} // namespace haar_lib
#line 4 "Mylib/Graph/TreeUtils/rerooting.cpp"
namespace haar_lib {
namespace rerooting_impl {
template <typename T, typename U, typename Merge, typename EdgeF, typename VertexF>
T rec1(
tree<U> &tr,
T id,
const Merge &merge,
const EdgeF &f,
const VertexF &g,
std::vector<std::vector<T>> &dp,
int cur,
int par = -1) {
T acc = id;
for (int i = 0; i < (int) tr[cur].size(); ++i) {
auto &e = tr[cur][i];
if (e.to == par) continue;
dp[cur][i] = rec1(tr, id, merge, f, g, dp, e.to, cur);
acc = merge(acc, f(dp[cur][i], e));
}
return g(acc, cur);
}
template <typename T, typename U, typename Merge, typename EdgeF, typename VertexF>
void rec2(
const tree<U> &tr,
T id,
const Merge &merge,
const EdgeF &f,
const VertexF &g,
std::vector<std::vector<T>> &dp,
int cur,
int par,
T value) {
const int l = tr[cur].size();
for (int i = 0; i < l; ++i) {
if (tr[cur][i].to == par) {
dp[cur][i] = value;
}
}
std::vector<T> left(l + 1, id), right(l + 1, id);
for (int i = 0; i < l - 1; ++i) {
const auto &e = tr[cur][i];
left[i + 1] = merge(left[i], f(dp[cur][i], e));
}
for (int i = l - 1; i >= 1; --i) {
const auto &e = tr[cur][i];
right[i - 1] = merge(right[i], f(dp[cur][i], e));
}
for (int i = 0; i < l; ++i) {
const auto &e = tr[cur][i];
if (e.to == par) continue;
rec2(tr, id, merge, f, g, dp, e.to, cur, g(merge(left[i], right[i]), cur));
}
}
} // namespace rerooting_impl
template <typename T, typename U, typename Merge, typename EdgeF, typename VertexF>
auto rerooting(tree<U> tr, T id, Merge merge, EdgeF f, VertexF g) {
const int N = tr.size();
std::vector<std::vector<T>> dp(N);
std::vector<T> ret(N, id);
for (int i = 0; i < N; ++i) dp[i].assign(tr[i].size(), id);
rerooting_impl::rec1(tr, id, merge, f, g, dp, 0);
rerooting_impl::rec2(tr, id, merge, f, g, dp, 0, -1, id);
for (int i = 0; i < N; ++i) {
for (int j = 0; j < (int) tr[i].size(); ++j) {
ret[i] = merge(ret[i], f(dp[i][j], tr[i][j]));
}
ret[i] = g(ret[i], i);
}
return ret;
}
} // namespace haar_lib
Rerooting DP
test/aoj/1595/main.test.cpp
test/yukicoder/922/main.test.cpp