This documentation is automatically generated by online-judge-tools/verification-helper
#include "tree/guni.hpp"
頂点数 $n$ の根付き木について,頂点集合の部分集合の列 $S = (S_0, \ldots, S_m)$ で,以下を満たすものを構築する.
本コードはこのような $S$ の一つを implicit に構築する.根付き木の全ての部分木に関する特定の問題を全ての部分木に対して解く必要があるような状況で,効率的な実装に役立つ.
木の上で DFS し,行きがけに頂点追加・帰りがけに頂点削除を行うのが基本方針だが,一部の頂点には複数回訪問する点で通常の DFS とは異なる.
Heavy-light decomposition を行う.ある頂点 $v$ に到達した際,以下の手続きを行う:
このアルゴリズムは以下の背景により $O(n \log n)$ である:
本テクニックを使用するまでもない例だが,根付き木の全ての部分木について「部分木を構成する頂点の id
の 2 乗和」は以下のように実装できる.
int N, root;
vector<pair<int, int>> edges(N - 1);
guni g(N);
for (auto [u, v] : edges) g.add_bi_edge(u, v);
std::vector<long long> ret(N);
long long sum_of_i_quads = 0;
auto Add = [&](int i) { sum_of_i_quads += (long long)i * i; };
auto Remove = [&](int i) { sum_of_i_quads -= (long long)i * i; };
auto Solve = [&](int i) { ret.at(i) = sum_of_i_quads; };
g.run(0, Add, Remove, Solve);
#pragma once
#include <vector>
// Guni / Sack / DSU on tree
// https://codeforces.com/blog/entry/44351
struct guni {
int n;
int last_root;
std::vector<std::vector<int>> to;
std::vector<int> sz, ver, st, ft; // subtree size / dfs order / subtree start / subtree fin
guni(int n_ = 0) : n(n_), last_root(-1), to(n_) {}
void add_bi_edge(int u, int v) { to.at(u).push_back(v), to.at(v).push_back(u); }
void sdfs(int v, int p) { // Build sz / ver / st / ft
st.at(v) = ver.size(), ver.push_back(v);
for (int u : to.at(v)) sz.at(v) += (u != p) ? (sdfs(u, v), sz.at(u)) : 0;
ft.at(v) = ver.size();
}
template <class F1, class F2, class F3>
void dfs(int v, int p, bool keep, F1 Add, F2 Remove, F3 Solve) {
int mx = -1, big_child = -1;
for (int u : to.at(v)) {
if (u != p and sz.at(u) > mx) mx = sz.at(u), big_child = u;
}
for (int u : to.at(v)) {
if (u != p and u != big_child) dfs(u, v, false, Add, Remove, Solve);
}
if (big_child != -1) dfs(big_child, v, true, Add, Remove, Solve);
for (int u : to.at(v)) {
if (u != p and u != big_child) {
for (int i = st.at(u); i < ft.at(u); ++i) Add(ver.at(i));
}
}
Add(v);
Solve(v);
if (!keep) {
for (int i = st.at(v); i < ft.at(v); ++i) Remove(ver.at(i));
}
}
template <class F1, class F2, class F3> void run(const int root, F1 Add, F2 Remove, F3 Solve) {
if (last_root != root) {
last_root = root, ver.clear(), st.resize(n), ft.resize(n), sz.assign(n, 1);
sdfs(root, -1);
}
dfs(root, -1, false, Add, Remove, Solve);
}
};
#line 2 "tree/guni.hpp"
#include <vector>
// Guni / Sack / DSU on tree
// https://codeforces.com/blog/entry/44351
struct guni {
int n;
int last_root;
std::vector<std::vector<int>> to;
std::vector<int> sz, ver, st, ft; // subtree size / dfs order / subtree start / subtree fin
guni(int n_ = 0) : n(n_), last_root(-1), to(n_) {}
void add_bi_edge(int u, int v) { to.at(u).push_back(v), to.at(v).push_back(u); }
void sdfs(int v, int p) { // Build sz / ver / st / ft
st.at(v) = ver.size(), ver.push_back(v);
for (int u : to.at(v)) sz.at(v) += (u != p) ? (sdfs(u, v), sz.at(u)) : 0;
ft.at(v) = ver.size();
}
template <class F1, class F2, class F3>
void dfs(int v, int p, bool keep, F1 Add, F2 Remove, F3 Solve) {
int mx = -1, big_child = -1;
for (int u : to.at(v)) {
if (u != p and sz.at(u) > mx) mx = sz.at(u), big_child = u;
}
for (int u : to.at(v)) {
if (u != p and u != big_child) dfs(u, v, false, Add, Remove, Solve);
}
if (big_child != -1) dfs(big_child, v, true, Add, Remove, Solve);
for (int u : to.at(v)) {
if (u != p and u != big_child) {
for (int i = st.at(u); i < ft.at(u); ++i) Add(ver.at(i));
}
}
Add(v);
Solve(v);
if (!keep) {
for (int i = st.at(v); i < ft.at(v); ++i) Remove(ver.at(i));
}
}
template <class F1, class F2, class F3> void run(const int root, F1 Add, F2 Remove, F3 Solve) {
if (last_root != root) {
last_root = root, ver.clear(), st.resize(n), ft.resize(n), sz.assign(n, 1);
sdfs(root, -1);
}
dfs(root, -1, false, Add, Remove, Solve);
}
};