virtual_tree.hpp¶
Vertices and parent edges of a compressed virtual tree.
把一次询问涉及的若干关键点及其 LCA 压成小树；适合多组树上关键点 DP。
Implementation¶
#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> vertices;
std::vector<int> parent;
std::vector<int> edge_length;
std::vector<bool> is_key;
};
/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
inline virtual_tree_result build_virtual_tree(const hld &tree,
std::vector<int> keys) {
virtual_tree_result result;
if (keys.empty()) {
return result;
}
for (int vertex : keys) {
assert(0 <= vertex && vertex < tree.n);
}
std::sort(keys.begin(), keys.end());
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
std::unordered_set<int> key_set(keys.begin(), keys.end());
std::sort(keys.begin(), keys.end(), [&](int first, int second) {
return tree.dfn[first] < tree.dfn[second];
});
result.vertices = keys;
for (int i = 1; i < int(keys.size()); i++) {
result.vertices.push_back(tree.lca(keys[i - 1], keys[i]));
}
std::sort(result.vertices.begin(), result.vertices.end(),
[&](int first, int second) {
return tree.dfn[first] < tree.dfn[second];
});
result.vertices.erase(
std::unique(result.vertices.begin(), result.vertices.end()),
result.vertices.end());
int size = int(result.vertices.size());
result.parent.assign(size, -1);
result.edge_length.assign(size, 0);
result.is_key.resize(size);
std::vector<int> stack;
for (int index = 0; index < size; index++) {
int vertex = result.vertices[index];
result.is_key[index] = key_set.contains(vertex);
while (!stack.empty() &&
!(tree.dfn[result.vertices[stack.back()]] <= tree.dfn[vertex] &&
tree.dfn[vertex] < tree.dfn[result.vertices[stack.back()]] +
tree.siz[result.vertices[stack.back()]])) {
stack.pop_back();
}
if (!stack.empty()) {
result.parent[index] = stack.back();
result.edge_length[index] =
tree.d[vertex] - tree.d[result.vertices[stack.back()]];
}
stack.push_back(index);
}
return result;
}
} // 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> tour_list;
void dfs2(int u, int t) {
top[u] = t, dfn[u] = idx++;
tour_list.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 goal = d[a] - k;
while (d[top[a]] > goal)
a = fa[top[a]];
int pos = dfn[a] - (d[a] - goal);
return tour_list[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 first_half = d[a] - d[anc];
int second_half = d[b] - d[anc];
if (k < 0 || k > first_half + second_half)
return -1;
if (k < first_half)
return get_kth_ancestor(a, k);
else
return get_kth_ancestor(b, first_half + second_half - k);
}
hld(const std::vector<std::vector<int>> &g = {}, const int &root = 0) {
if (!g.empty())
build(g, root);
}
void build(const std::vector<std::vector<int>> &g = {}, const int &root = 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[root] = 1;
dfs(root, -1);
idx = 0;
tour_list.clear();
dfs2(root, root);
}
/// @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> vertices;
std::vector<int> parent;
std::vector<int> edge_length;
std::vector<bool> is_key;
};
/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
inline virtual_tree_result build_virtual_tree(const hld &tree,
std::vector<int> keys) {
virtual_tree_result result;
if (keys.empty()) {
return result;
}
for (int vertex : keys) {
assert(0 <= vertex && vertex < tree.n);
}
std::sort(keys.begin(), keys.end());
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
std::unordered_set<int> key_set(keys.begin(), keys.end());
std::sort(keys.begin(), keys.end(), [&](int first, int second) {
return tree.dfn[first] < tree.dfn[second];
});
result.vertices = keys;
for (int i = 1; i < int(keys.size()); i++) {
result.vertices.push_back(tree.lca(keys[i - 1], keys[i]));
}
std::sort(result.vertices.begin(), result.vertices.end(),
[&](int first, int second) {
return tree.dfn[first] < tree.dfn[second];
});
result.vertices.erase(
std::unique(result.vertices.begin(), result.vertices.end()),
result.vertices.end());
int size = int(result.vertices.size());
result.parent.assign(size, -1);
result.edge_length.assign(size, 0);
result.is_key.resize(size);
std::vector<int> stack;
for (int index = 0; index < size; index++) {
int vertex = result.vertices[index];
result.is_key[index] = key_set.contains(vertex);
while (!stack.empty() &&
!(tree.dfn[result.vertices[stack.back()]] <= tree.dfn[vertex] &&
tree.dfn[vertex] < tree.dfn[result.vertices[stack.back()]] +
tree.siz[result.vertices[stack.back()]])) {
stack.pop_back();
}
if (!stack.empty()) {
result.parent[index] = stack.back();
result.edge_length[index] =
tree.d[vertex] - tree.d[result.vertices[stack.back()]];
}
stack.push_back(index);
}
return result;
}
} // namespace noya