This documentation is automatically generated by online-judge-tools/verification-helper
View the Project on GitHub kmyk/competitive-programming-library
#include "utils/dsu_on_tree.hpp"
#pragma once
#include <functional>
#include <stack>
#include <vector>
#include "../graph/subtree.hpp"
#include "../utils/macros.hpp"
/**
* @brief DSU on tree (sack)
* @arg g is a tree
* @arg root
* @arg add is a function object which takes a index of a vertex
* @arg sub is a function object which takes a index of a vertex
* @arg callback is a function object which takes a index of a vertex
* @note for each x, add(x) and sub(x) are called O(log n) times
* @note O(n log n) if add, sub, and callback are O(1)
* @see https://codeforces.com/blog/entry/44351
* @note sub(x) can be implemented as reset(), because sub(x) is called until it becomes empty after sub(x) is called once
*/
template <class Add, class Sub, class Callback>
void dsu_on_tree(const std::vector<std::vector<int> > & g, int root, Add & add, Sub & sub, Callback & callback) {
auto info = prepare_subtree_info(g, root);
auto subtree_apply = [&](int x, auto & f) {
std::stack<int> stk;
stk.push(x);
while (not stk.empty()) {
int y = stk.top();
stk.pop();
f(y);
for (int z : g[y]) if (z != info[y].parent) {
stk.push(z);
}
}
};
std::function<void (int, bool)> go = [&](int x, bool keep) {
// leaf
if (info[x].size == 1) {
add(x);
callback(x);
if (not keep) {
sub(x);
}
return;
}
// choose the heavy child
int z = *max_element(ALL(g[x]), [&](int y1, int y2) {
int size1 = (y1 == info[x].parent ? -1 : info[y1].size);
int size2 = (y2 == info[x].parent ? -1 : info[y2].size);
return size1 < size2;
});
// go light
for (int y : g[x]) if (y != info[x].parent) {
if (y != z) {
go(y, false);
}
}
// go heavy
go(z, true);
for (int y : g[x]) if (y != info[x].parent) {
if (y != z) {
subtree_apply(y, add);
}
}
add(x);
callback(x);
if (not keep) {
subtree_apply(x, sub);
}
};
go(root, false);
}
#line 2 "utils/dsu_on_tree.hpp"
#include <functional>
#include <stack>
#include <vector>
#line 2 "graph/subtree.hpp"
#include <algorithm>
#line 4 "graph/subtree.hpp"
struct subtree_info_t {
int parent; // in the entire tree
int depth; // in the entire tree
int size; // of the subtree
int height; // of the subtree
};
/**
* @brief subtree info / それぞれの部分木の size とか height とかをまとめて求めておいてくれるやつ
* @arg g must be a tree
* @note O(n) time
* @note O(n) space on heap
*/
std::vector<subtree_info_t> prepare_subtree_info(std::vector<std::vector<int> > const & g, int root) {
int n = g.size();
std::vector<subtree_info_t> info(n, (subtree_info_t) { -1, -1, -1, -1 });
std::vector<int> topological(n);
topological[0] = root;
info[root].parent = root;
info[root].depth = 0;
int r = 1;
for (int l = 0; l < r; ++ l) {
int i = topological[l];
for (int j : g[i]) if (j != info[i].parent) {
topological[r ++] = j;
info[j].parent = i;
info[j].depth = info[i].depth + 1;
}
}
while ((-- r) >= 0) {
int i = topological[r];
info[i].size = 1;
info[i].height = 0;
for (int j : g[i]) if (j != info[i].parent) {
info[i].size += info[j].size;
info[i].height = std::max(info[i].height, info[j].height + 1);
}
}
info[root].parent = -1;
return info;
}
#line 2 "utils/macros.hpp"
#define REP(i, n) for (int i = 0; (i) < (int)(n); ++ (i))
#define REP3(i, m, n) for (int i = (m); (i) < (int)(n); ++ (i))
#define REP_R(i, n) for (int i = (int)(n) - 1; (i) >= 0; -- (i))
#define REP3R(i, m, n) for (int i = (int)(n) - 1; (i) >= (int)(m); -- (i))
#define ALL(x) std::begin(x), std::end(x)
#line 7 "utils/dsu_on_tree.hpp"
/**
* @brief DSU on tree (sack)
* @arg g is a tree
* @arg root
* @arg add is a function object which takes a index of a vertex
* @arg sub is a function object which takes a index of a vertex
* @arg callback is a function object which takes a index of a vertex
* @note for each x, add(x) and sub(x) are called O(log n) times
* @note O(n log n) if add, sub, and callback are O(1)
* @see https://codeforces.com/blog/entry/44351
* @note sub(x) can be implemented as reset(), because sub(x) is called until it becomes empty after sub(x) is called once
*/
template <class Add, class Sub, class Callback>
void dsu_on_tree(const std::vector<std::vector<int> > & g, int root, Add & add, Sub & sub, Callback & callback) {
auto info = prepare_subtree_info(g, root);
auto subtree_apply = [&](int x, auto & f) {
std::stack<int> stk;
stk.push(x);
while (not stk.empty()) {
int y = stk.top();
stk.pop();
f(y);
for (int z : g[y]) if (z != info[y].parent) {
stk.push(z);
}
}
};
std::function<void (int, bool)> go = [&](int x, bool keep) {
// leaf
if (info[x].size == 1) {
add(x);
callback(x);
if (not keep) {
sub(x);
}
return;
}
// choose the heavy child
int z = *max_element(ALL(g[x]), [&](int y1, int y2) {
int size1 = (y1 == info[x].parent ? -1 : info[y1].size);
int size2 = (y2 == info[x].parent ? -1 : info[y2].size);
return size1 < size2;
});
// go light
for (int y : g[x]) if (y != info[x].parent) {
if (y != z) {
go(y, false);
}
}
// go heavy
go(z, true);
for (int y : g[x]) if (y != info[x].parent) {
if (y != z) {
subtree_apply(y, add);
}
}
add(x);
callback(x);
if (not keep) {
subtree_apply(x, sub);
}
};
go(root, false);
}