virtual_tree.hpp¶
把一次询问涉及的若干关键点及其 LCA 压成小树；适合多组树上关键点 DP。
Complexity: Time: O(k log k) for k key vertices after LCA preprocessing. Space: O(k).
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(k log k) for k key vertices after LCA preprocessing.
/// Space: O(k).
#include "noya/heavy_light_decomposition.hpp"
#include <algorithm>
#include <cassert>
#include <unordered_set>
#include <vector>
namespace noya {
/// @brief Vertices and parent edges of a compressed virtual tree.
struct virtual_tree_result {
std::vector<int> vs;
std::vector<int> fa;
std::vector<int> len;
std::vector<bool> key;
};
/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
inline virtual_tree_result build_virtual_tree(const hld &tr,
std::vector<int> vs1) {
virtual_tree_result res;
if (vs1.empty()) {
return res;
}
for (int u : vs1) {
assert(0 <= u && u < tr.n);
}
std::sort(vs1.begin(), vs1.end());
vs1.erase(std::unique(vs1.begin(), vs1.end()), vs1.end());
std::unordered_set<int> ks(vs1.begin(), vs1.end());
std::sort(vs1.begin(), vs1.end(),
[&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
res.vs = vs1;
for (int i = 1; i < int(vs1.size()); i++) {
res.vs.push_back(tr.lca(vs1[i - 1], vs1[i]));
}
std::sort(res.vs.begin(), res.vs.end(),
[&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
res.vs.erase(std::unique(res.vs.begin(), res.vs.end()), res.vs.end());
int sz = int(res.vs.size());
res.fa.assign(sz, -1);
res.len.assign(sz, 0);
res.key.resize(sz);
std::vector<int> stk;
for (int idx = 0; idx < sz; idx++) {
int u = res.vs[idx];
res.key[idx] = ks.contains(u);
while (!stk.empty() && !(tr.dfn[res.vs[stk.back()]] <= tr.dfn[u] &&
tr.dfn[u] < tr.dfn[res.vs[stk.back()]] +
tr.siz[res.vs[stk.back()]])) {
stk.pop_back();
}
if (!stk.empty()) {
res.fa[idx] = stk.back();
res.len[idx] = tr.d[u] - tr.d[res.vs[stk.back()]];
}
stk.push_back(idx);
}
return res;
}
} // namespace noya
#ifndef NOYA_VIRTUAL_TREE_HPP
#define NOYA_VIRTUAL_TREE_HPP 1
/// @complexity Time: O(k log k) for k key vertices after LCA preprocessing.
/// Space: O(k).
#include "noya/heavy_light_decomposition.hpp"
#include <algorithm>
#include <cassert>
#include <unordered_set>
#include <vector>
namespace noya {
/// @brief Vertices and parent edges of a compressed virtual tree.
struct virtual_tree_result {
std::vector<int> vs;
std::vector<int> fa;
std::vector<int> len;
std::vector<bool> key;
};
/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
inline virtual_tree_result build_virtual_tree(const hld &tr,
std::vector<int> vs1) {
virtual_tree_result res;
if (vs1.empty()) {
return res;
}
for (int u : vs1) {
assert(0 <= u && u < tr.n);
}
std::sort(vs1.begin(), vs1.end());
vs1.erase(std::unique(vs1.begin(), vs1.end()), vs1.end());
std::unordered_set<int> ks(vs1.begin(), vs1.end());
std::sort(vs1.begin(), vs1.end(),
[&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
res.vs = vs1;
for (int i = 1; i < int(vs1.size()); i++) {
res.vs.push_back(tr.lca(vs1[i - 1], vs1[i]));
}
std::sort(res.vs.begin(), res.vs.end(),
[&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
res.vs.erase(std::unique(res.vs.begin(), res.vs.end()), res.vs.end());
int sz = int(res.vs.size());
res.fa.assign(sz, -1);
res.len.assign(sz, 0);
res.key.resize(sz);
std::vector<int> stk;
for (int idx = 0; idx < sz; idx++) {
int u = res.vs[idx];
res.key[idx] = ks.contains(u);
while (!stk.empty() && !(tr.dfn[res.vs[stk.back()]] <= tr.dfn[u] &&
tr.dfn[u] < tr.dfn[res.vs[stk.back()]] +
tr.siz[res.vs[stk.back()]])) {
stk.pop_back();
}
if (!stk.empty()) {
res.fa[idx] = stk.back();
res.len[idx] = tr.d[u] - tr.d[res.vs[stk.back()]];
}
stk.push_back(idx);
}
return res;
}
} // namespace noya
#endif // NOYA_VIRTUAL_TREE_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <tuple>
#include <unordered_set>
#include <utility>
#include <vector>
/// @complexity Time: O(k log k) for k key vertices after LCA preprocessing.
/// Space: O(k).
/// @complexity Time: O(n) build, O(log n) LCA/path decomposition, O(1) subtree interval.
/// Space: O(n).
namespace noya {
/// @brief Heavy-light decomposition for path and subtree queries on trees.
struct hld {
std::vector<std::vector<int>> G;
int n;
std::vector<int> dfn, siz, son, top, d, fa;
int idx;
void dfs(int u, int p) {
siz[u] = 1;
for (auto v : G[u])
if (v != p) {
fa[v] = u;
d[v] = d[u] + 1;
dfs(v, u), siz[u] += siz[v];
if (son[u] == -1 || siz[v] > siz[son[u]])
son[u] = v;
}
}
std::vector<int> tl;
void dfs2(int u, int t) {
top[u] = t, dfn[u] = idx++;
tl.push_back(u);
if (son[u] != -1)
dfs2(son[u], t);
for (auto v : G[u])
if (top[v] == -1)
dfs2(v, v);
}
/// @brief Return the k-th ancestor of node a, or -1 if k > depth(a).
int get_kth_ancestor(int a, int k) const {
if (k < 0 || k >= d[a])
return -1;
int dst = d[a] - k;
while (d[top[a]] > dst)
a = fa[top[a]];
int pos = dfn[a] - (d[a] - dst);
return tl[pos];
}
/// @brief Return the k-th node (0-indexed) on the path from a to b, or -1 if out of range.
int get_kth_node_on_path(int a, int b, int k) const {
int anc = lca(a, b);
int lhs = d[a] - d[anc];
int rhs = d[b] - d[anc];
if (k < 0 || k > lhs + rhs)
return -1;
if (k < lhs)
return get_kth_ancestor(a, k);
else
return get_kth_ancestor(b, lhs + rhs - k);
}
hld(const std::vector<std::vector<int>> &g = {}, const int &rt = 0) {
if (!g.empty())
build(g, rt);
}
void build(const std::vector<std::vector<int>> &g = {}, const int &rt = 0) {
n = g.size();
G = g;
siz.assign(n, 0);
dfn.assign(n, -1);
son.assign(n, -1);
top.assign(n, -1);
d.assign(n, 0);
fa.assign(n, -1);
d[rt] = 1;
dfs(rt, -1);
idx = 0;
tl.clear();
dfs2(rt, rt);
}
/// @brief Check if a is in the subtree of b.
bool is_subtree(int a, int b) const {
if (dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b]) {
return true;
} else {
return false;
}
}
/// @brief Return the lowest common ancestor of x and y.
int lca(int x, int y) const {
while (top[x] != top[y]) {
if (d[top[x]] < d[top[y]])
std::swap(x, y);
x = fa[top[x]];
}
return d[x] < d[y] ? x : y;
}
/// @brief Decompose path x->y into chain segments. @return (dfn_l, dfn_r, direction).
std::vector<std::tuple<int, int, bool>> chain(int x, int y) const {
assert(0 <= x && x < n);
assert(0 <= y && y < n);
std::vector<std::tuple<int, int, bool>> L, R;
while (top[x] != top[y]) {
assert(0 <= x && x < n);
assert(0 <= y && y < n);
if (d[top[x]] > d[top[y]]) {
L.emplace_back(dfn[top[x]], dfn[x] + 1, false);
x = fa[top[x]];
} else {
R.emplace_back(dfn[top[y]], dfn[y] + 1, true);
y = fa[top[y]];
}
}
if (dfn[y] < dfn[x])
L.emplace_back(dfn[y], dfn[x] + 1, false);
else
R.emplace_back(dfn[x], dfn[y] + 1, true);
reverse(R.begin(), R.end());
L.insert(L.end(), R.begin(), R.end());
return L;
}
/// @brief Return the DFN range [l, r) for the subtree of node a.
std::array<int, 2> subtree(int a) const { return {dfn[a], dfn[a] + siz[a]}; }
/// @brief Return the LCA when the tree is re-rooted at c.
int rooted_lca(int a, int b, int c) const {
return lca(a, b) ^ lca(a, c) ^ lca(b, c);
}
/// @brief Compute the intersection of paths (a,b) and (c,d) as a pair of endpoints.
std::pair<int, int> intersection(int a, int b, int c, int d) const {
int ab = lca(a, b), ac = lca(a, c), ad = lca(a, d);
int bc = lca(b, c), bd = lca(b, d), cd = lca(c, d);
int x = ab ^ ac ^ bc;
int y = ab ^ ad ^ bd;
if (x != y) {
return {x, y};
}
int z = ac ^ ad ^ cd;
if (x != z) {
x = -1;
}
return {x, x};
}
std::pair<int, int> intersection(std::pair<int, int> a,
std::pair<int, int> b) const {
return intersection(a.first, a.second, b.first, b.second);
}
};
} // namespace noya
namespace noya {
/// @brief Vertices and parent edges of a compressed virtual tree.
struct virtual_tree_result {
std::vector<int> vs;
std::vector<int> fa;
std::vector<int> len;
std::vector<bool> key;
};
/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
inline virtual_tree_result build_virtual_tree(const hld &tr,
std::vector<int> vs1) {
virtual_tree_result res;
if (vs1.empty()) {
return res;
}
for (int u : vs1) {
assert(0 <= u && u < tr.n);
}
std::sort(vs1.begin(), vs1.end());
vs1.erase(std::unique(vs1.begin(), vs1.end()), vs1.end());
std::unordered_set<int> ks(vs1.begin(), vs1.end());
std::sort(vs1.begin(), vs1.end(),
[&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
res.vs = vs1;
for (int i = 1; i < int(vs1.size()); i++) {
res.vs.push_back(tr.lca(vs1[i - 1], vs1[i]));
}
std::sort(res.vs.begin(), res.vs.end(),
[&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
res.vs.erase(std::unique(res.vs.begin(), res.vs.end()), res.vs.end());
int sz = int(res.vs.size());
res.fa.assign(sz, -1);
res.len.assign(sz, 0);
res.key.resize(sz);
std::vector<int> stk;
for (int idx = 0; idx < sz; idx++) {
int u = res.vs[idx];
res.key[idx] = ks.contains(u);
while (!stk.empty() && !(tr.dfn[res.vs[stk.back()]] <= tr.dfn[u] &&
tr.dfn[u] < tr.dfn[res.vs[stk.back()]] +
tr.siz[res.vs[stk.back()]])) {
stk.pop_back();
}
if (!stk.empty()) {
res.fa[idx] = stk.back();
res.len[idx] = tr.d[u] - tr.d[res.vs[stk.back()]];
}
stk.push_back(idx);
}
return res;
}
} // namespace noya