競技プログラミング用 C++ ライブラリ
tree/hld.hpp · GitHub で見る
共通ページ
#include "lib/tree/hld.hpp"
HLD tree(3);
tree.add_edge(0, 1);
tree.add_edge(1, 2);
tree.build();
auto segments = tree.path(2, 0);
auto [l, r] = tree.subtree(1);
| 操作 | 内容・返り値 | 計算量 |
|---|---|---|
HLD(n = 0) |
頂点数を指定 | \(O(n)\) |
add_edge(u,v) |
無向辺を追加し、辺 ID を返す | 償却 \(O(1)\) |
build(root = 0) |
指定した根で木を分解 | \(O(n)\) |
build_from_parents(parents) |
親配列から木を構築し、分解 | \(O(n)\) |
lca(u,v) / dist(u,v)
|
最小共通祖先/経路の辺数 | \(O(\log n)\) |
is_ancestor(u,v) |
u が v の祖先か。自身も含む |
\(O(1)\) |
ancestor(v,k) |
k 個上の祖先。範囲外は -1 |
\(O(\log n)\) |
jump(u,v,k) |
u を 0 番目とする経路上の頂点。範囲外は -1 |
\(O(\log n)\) |
subtree(v,edge = false) |
部分木の添字区間。辺の場合は親への辺を除く | \(O(1)\) |
edge_vertex(id) |
辺 ID の深い側の頂点 | \(O(1)\) |
path(u,v,edge = false) |
経路順の Segment{l,r,reversed} の列 |
\(O(\log n)\) |
for_each_path(u,v,f,edge = false) |
経路順に f(l,r,reversed) を呼ぶ |
\(O(\log n)\) |
parent[v] / depth[v] / size[v]
|
親/深さ/部分木の頂点数 | \(O(1)\) |
heavy[v] / head[v]
|
重い子/連鎖の先頭 | \(O(1)\) |
in[v] / out[v] / order[i]
|
部分木の左端/右端/添字に対応する頂点 | \(O(1)\) |
n / edges[id]
|
全頂点数/辺の両端 | \(O(1)\) |
無向木を扱い、build(root = 0) で根を指定します。頂点値は Segtree などの in[v] に置きます。
親配列がある場合は build_from_parents(std::vector<int> parents) を使えます。頂点数は配列の長さとなり、根の親を -1、それ以外の親を頂点番号で指定します。根はちょうど 1 頂点で、全頂点が根につながる木である必要があります。頂点番号は任意の順序で構いません。空配列も渡せます。
#include "lib/tree/hld.hpp"
HLD rooted;
rooted.build_from_parents({-1, 0, 1});
auto root = rooted.parent[0]; // -1
auto distance = rooted.dist(0, 2); // 2
既存の頂点・辺は置き換わります。edges には根以外の頂点 v の (parents[v], v) が v の昇順で入り、その順番が辺 ID になります。配列を以降使わない場合は build_from_parents(std::move(parents)) と渡せます。
in[v] は木を分解した後の配列上の位置で、元の頂点番号 v とは別です。元の頂点値が value[v] なら、長さ n の配列に a[tree.in[v]] = value[v] と並べ替えてから Segtree を構築します。subtree(v) の [l,r) と path が返す区間は、この配列に対する範囲です。tree.order[i] で位置 i に対応する元の頂点番号へ戻せます。
path(u,v,edge = false) は経路を u から v の順に並べた Segment{l,r,reversed} の列を返します。各区間は \([l,r)\) で、reversed が真なら降順に読みます。非可換な演算では、区間の正順・逆順の積を用意して使い分けます。for_each_path(u,v,f,edge = false) は同じ順序で f(l,r,reversed) を呼び、区間列のメモリ確保を省きます。
次は in[v] 順の値で構築した区間和のセグメント木 seg を使う部分です。u・v は経路の両端で、和は可換なので reversed を使いません。
long long answer = 0;
tree.for_each_path(u, v, [&](int l, int r, bool) { answer += seg.prod(l, r); });
辺値は in[tree.edge_vertex(id)] に置き、edge = true を指定します。id は add_edge の返り値です。subtree(v,edge = false) は部分木の区間を返し、辺の場合は v と親の辺を除きます。
lca(u,v)・dist(u,v)・is_ancestor(u,v)、ancestor(v,k)(k 個上の祖先)・jump(u,v,k)(経路の 0 始まり k 番目の頂点)が使えます。ancestor・jump は範囲外なら -1 です。
構築後の各配列は参照用です。parent[v] は親(根は -1)、depth[v] は深さ、size[v] は部分木サイズ、heavy[v] は重い子(なければ -1)、head[v] は連鎖の先頭です。in[v]・out[v] は部分木の区間、order[i] は添字 i の頂点です。
構築・空間は \(O(n)\)、lca・dist・ancestor・jump・パス分解は \(O(\log n)\)、それ以外のクエリは \(O(1)\) です。分解される区間数は \(O(\log n)\)、Segtree と組み合わせたパスクエリは \(O(\log^2 n)\) です。
#ifndef HLD_HPP
#define HLD_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <cstddef>
#include <limits>
#include <utility>
#include <vector>
class HLD {
int shape_ = 0, center_ = 0;
struct Node { int position, head, parent, offset; };
std::vector<Node> nodes;
void check([[maybe_unused]] int v) const { assert(0 <= v && v < n && int(in.size()) == n); }
void initialize() {
depth.assign(n, 0);
size.assign(n, 1);
heavy.assign(n, -1);
head.resize(n), in.resize(n), out.resize(n), order.resize(n);
nodes.resize(n);
shape_ = center_ = 0;
}
void finish(int root) {
in[root] = 0, out[root] = n, head[root] = root;
for (int i = n - 2; i >= 0; --i) {
int v = order[i], p = parent[v];
depth[v] = depth[p] + 1;
if (heavy[p] == v) in[v] = in[p] + 1, head[v] = head[p];
else in[v] = out[p] -= size[v], head[v] = v;
out[v] = in[v] + size[v];
}
for (int v = 0; v < n; ++v) out[v] = in[v] + size[v], order[in[v]] = v;
parent[root] = -1;
for (int v = 0; v < n; ++v) nodes[v] = {in[v], in[head[v]], parent[head[v]], depth[v] - in[v]};
}
public:
struct Segment {
int l, r;
bool reversed;
};
int n;
std::vector<std::pair<int, int>> edges;
std::vector<int> parent, depth, size, heavy, head, in, out, order;
explicit HLD(int n = 0) : n(n) {
assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
edges.reserve(std::max(0, n - 1));
}
int add_edge(int u, int v) {
assert(0 <= u && u < n && 0 <= v && v < n);
edges.emplace_back(u, v);
return int(edges.size()) - 1;
}
void build(int root = 0) {
assert(n >= 0 && edges.size() == std::size_t(std::max(0, n - 1)));
assert(!n || (0 <= root && root < n));
parent.assign(n, -1);
initialize();
if (!n) return;
std::vector<int> degree(n), neighbors(n);
for (auto [u, v] : edges) {
++degree[u], ++degree[v];
neighbors[u] ^= v, neighbors[v] ^= u;
}
center_ = int(std::max_element(degree.begin(), degree.end()) - degree.begin());
shape_ = degree[center_] <= 2 ? 1 : degree[center_] == n - 1 ? 2 : 0;
degree[root] += 2;
int count = 0;
for (int v = 0; v < n; ++v) if (degree[v] == 1) order[count++] = v;
for (int i = 0; i < count; ++i) {
int v = order[i], p = neighbors[v];
parent[v] = p;
size[p] += size[v];
if (heavy[p] == -1 || size[heavy[p]] < size[v]) heavy[p] = v;
neighbors[p] ^= v;
if (--degree[p] == 1) order[count++] = p;
}
assert(count == n - 1);
finish(root);
}
void build_from_parents(std::vector<int> parents) {
assert(parents.size() <= std::size_t(std::numeric_limits<int>::max() / 2));
n = static_cast<int>(parents.size());
parent = std::move(parents);
initialize();
edges.clear();
edges.reserve(std::max(0, n - 1));
if (!n) return;
std::fill(out.begin(), out.end(), 0);
int root = -1;
bool ordered = true;
for (int v = 0; v < n; ++v) {
int p = parent[v];
assert(-1 <= p && p < n && p != v);
if (p == -1) {
assert(root == -1);
root = v;
} else {
++out[p];
edges.emplace_back(p, v);
ordered &= p < v;
}
}
assert(root != -1);
int maximum = 0;
for (int v = 0; v < n; ++v) {
int degree = out[v] + (v != root);
if (degree > maximum) maximum = degree, center_ = v;
}
shape_ = maximum <= 2 ? 1 : maximum == n - 1 ? 2 : 0;
if (ordered) {
for (int v = n - 1; v > 0; --v) {
int p = parent[v];
size[p] += size[v];
if (heavy[p] == -1 || size[heavy[p]] < size[v]) heavy[p] = v;
order[n - 1 - v] = v;
}
} else {
int count = 0;
for (int v = 0; v < n; ++v) if (!out[v] && v != root) order[count++] = v;
for (int i = 0; i < count; ++i) {
int v = order[i], p = parent[v];
size[p] += size[v];
if (heavy[p] == -1 || size[heavy[p]] < size[v]) heavy[p] = v;
if (!--out[p] && p != root) order[count++] = p;
}
assert(count == n - 1);
}
finish(root);
}
int lca(int u, int v) const {
check(u), check(v);
if (shape_ == 1) return head[u] != head[v] ? order[0] : in[u] < in[v] ? u : v;
if (shape_ == 2) return u == v ? u : u == order[0] || v == order[0] ? order[0] : center_;
while (nodes[u].head != nodes[v].head) {
if (nodes[u].position > nodes[v].position) u = nodes[u].parent;
else v = nodes[v].parent;
}
return nodes[u].position < nodes[v].position ? u : v;
}
int dist(int u, int v) const {
int w = lca(u, v);
return depth[u] + depth[v] - 2 * depth[w];
}
bool is_ancestor(int u, int v) const {
check(u), check(v);
return in[u] <= in[v] && in[v] < out[u];
}
int ancestor(int v, int k) const {
check(v);
if (shape_) {
if (k < 0 || k > depth[v]) return -1;
if (!k) return v;
if (k == depth[v]) return order[0];
return shape_ == 2 ? parent[v] : order[in[v] - k];
}
if (k < 0) return -1;
int target = nodes[v].position + nodes[v].offset - k;
if (target < 0) return -1;
while (nodes[v].head + nodes[v].offset > target) v = nodes[v].parent;
return order[target - nodes[v].offset];
}
int jump(int u, int v, int k) const {
check(u), check(v);
if (k < 0) return -1;
if (!k) return u;
if (shape_ == 1) {
if (head[u] == head[v]) {
int delta = in[v] - in[u];
if (k > (delta < 0 ? -delta : delta)) return -1;
return order[in[u] + (delta < 0 ? -k : k)];
}
int a = depth[u], b = depth[v];
if (k > a + b) return -1;
if (k == a) return order[0];
return k < a ? order[in[u] - k] : order[in[v] - (a + b - k)];
}
if (shape_ == 2) {
int d = (u != center_) + (v != center_);
if (u == v || k > d) return -1;
return k == d ? v : center_;
}
int from = u, to = v, target = nodes[u].position + nodes[u].offset - k;
while (nodes[u].head != nodes[v].head) {
if (nodes[u].position > nodes[v].position) {
if (nodes[u].head + nodes[u].offset <= target) return order[target - nodes[u].offset];
u = nodes[u].parent;
} else v = nodes[v].parent;
}
int w = nodes[u].position < nodes[v].position ? u : v;
int d = nodes[w].position + nodes[w].offset;
int length = nodes[from].position + nodes[from].offset + nodes[to].position + nodes[to].offset - 2 * d;
if (k > length) return -1;
if (target >= d) return order[target - nodes[u].offset];
return ancestor(to, length - k);
}
std::pair<int, int> subtree(int v, bool edge = false) const {
check(v);
return {in[v] + edge, out[v]};
}
int edge_vertex(int id) const {
assert(0 <= id && id < int(edges.size()));
auto [u, v] = edges[id];
check(u);
return parent[u] == v ? u : v;
}
template <class F>
void for_each_path(int u, int v, F &&f, bool edge = false) const {
check(u), check(v);
if (shape_ == 1) {
int x = in[u], y = in[v];
if (head[u] == head[v]) {
int l = std::min(x, y) + edge, r = std::max(x, y) + 1;
if (l < r) f(l, r, x > y);
} else if (head[u] == order[0]) {
if (edge < x + 1) f(int(edge), x + 1, true);
f(in[head[v]], y + 1, false);
} else {
f(in[head[u]], x + 1, true);
if (edge < y + 1) f(int(edge), y + 1, false);
}
return;
}
std::array<std::array<int, 2>, 32> right;
int count = 0;
while (nodes[u].head != nodes[v].head) {
if (nodes[u].position > nodes[v].position) {
f(nodes[u].head, nodes[u].position + 1, true);
u = nodes[u].parent;
} else {
assert(count < int(right.size()));
right[count++] = {nodes[v].head, nodes[v].position + 1};
v = nodes[v].parent;
}
}
bool reversed = nodes[u].position > nodes[v].position;
int l = std::min(nodes[u].position, nodes[v].position) + edge, r = std::max(nodes[u].position, nodes[v].position) + 1;
if (l < r) f(l, r, reversed);
while (count) {
auto [a, b] = right[--count];
f(a, b, false);
}
}
std::vector<Segment> path(int u, int v, bool edge = false) const {
std::vector<Segment> result;
for_each_path(u, v, [&](int l, int r, bool reversed) { result.push_back({l, r, reversed}); }, edge);
return result;
}
};
#endif
#line 1 "tree/hld.hpp"
#include <algorithm>
#include <array>
#include <cassert>
#include <cstddef>
#include <limits>
#include <utility>
#include <vector>
class HLD {
int shape_ = 0, center_ = 0;
struct Node { int position, head, parent, offset; };
std::vector<Node> nodes;
void check([[maybe_unused]] int v) const { assert(0 <= v && v < n && int(in.size()) == n); }
void initialize() {
depth.assign(n, 0);
size.assign(n, 1);
heavy.assign(n, -1);
head.resize(n), in.resize(n), out.resize(n), order.resize(n);
nodes.resize(n);
shape_ = center_ = 0;
}
void finish(int root) {
in[root] = 0, out[root] = n, head[root] = root;
for (int i = n - 2; i >= 0; --i) {
int v = order[i], p = parent[v];
depth[v] = depth[p] + 1;
if (heavy[p] == v) in[v] = in[p] + 1, head[v] = head[p];
else in[v] = out[p] -= size[v], head[v] = v;
out[v] = in[v] + size[v];
}
for (int v = 0; v < n; ++v) out[v] = in[v] + size[v], order[in[v]] = v;
parent[root] = -1;
for (int v = 0; v < n; ++v) nodes[v] = {in[v], in[head[v]], parent[head[v]], depth[v] - in[v]};
}
public:
struct Segment {
int l, r;
bool reversed;
};
int n;
std::vector<std::pair<int, int>> edges;
std::vector<int> parent, depth, size, heavy, head, in, out, order;
explicit HLD(int n = 0) : n(n) {
assert(0 <= n && n <= std::numeric_limits<int>::max() / 2);
edges.reserve(std::max(0, n - 1));
}
int add_edge(int u, int v) {
assert(0 <= u && u < n && 0 <= v && v < n);
edges.emplace_back(u, v);
return int(edges.size()) - 1;
}
void build(int root = 0) {
assert(n >= 0 && edges.size() == std::size_t(std::max(0, n - 1)));
assert(!n || (0 <= root && root < n));
parent.assign(n, -1);
initialize();
if (!n) return;
std::vector<int> degree(n), neighbors(n);
for (auto [u, v] : edges) {
++degree[u], ++degree[v];
neighbors[u] ^= v, neighbors[v] ^= u;
}
center_ = int(std::max_element(degree.begin(), degree.end()) - degree.begin());
shape_ = degree[center_] <= 2 ? 1 : degree[center_] == n - 1 ? 2 : 0;
degree[root] += 2;
int count = 0;
for (int v = 0; v < n; ++v) if (degree[v] == 1) order[count++] = v;
for (int i = 0; i < count; ++i) {
int v = order[i], p = neighbors[v];
parent[v] = p;
size[p] += size[v];
if (heavy[p] == -1 || size[heavy[p]] < size[v]) heavy[p] = v;
neighbors[p] ^= v;
if (--degree[p] == 1) order[count++] = p;
}
assert(count == n - 1);
finish(root);
}
void build_from_parents(std::vector<int> parents) {
assert(parents.size() <= std::size_t(std::numeric_limits<int>::max() / 2));
n = static_cast<int>(parents.size());
parent = std::move(parents);
initialize();
edges.clear();
edges.reserve(std::max(0, n - 1));
if (!n) return;
std::fill(out.begin(), out.end(), 0);
int root = -1;
bool ordered = true;
for (int v = 0; v < n; ++v) {
int p = parent[v];
assert(-1 <= p && p < n && p != v);
if (p == -1) {
assert(root == -1);
root = v;
} else {
++out[p];
edges.emplace_back(p, v);
ordered &= p < v;
}
}
assert(root != -1);
int maximum = 0;
for (int v = 0; v < n; ++v) {
int degree = out[v] + (v != root);
if (degree > maximum) maximum = degree, center_ = v;
}
shape_ = maximum <= 2 ? 1 : maximum == n - 1 ? 2 : 0;
if (ordered) {
for (int v = n - 1; v > 0; --v) {
int p = parent[v];
size[p] += size[v];
if (heavy[p] == -1 || size[heavy[p]] < size[v]) heavy[p] = v;
order[n - 1 - v] = v;
}
} else {
int count = 0;
for (int v = 0; v < n; ++v) if (!out[v] && v != root) order[count++] = v;
for (int i = 0; i < count; ++i) {
int v = order[i], p = parent[v];
size[p] += size[v];
if (heavy[p] == -1 || size[heavy[p]] < size[v]) heavy[p] = v;
if (!--out[p] && p != root) order[count++] = p;
}
assert(count == n - 1);
}
finish(root);
}
int lca(int u, int v) const {
check(u), check(v);
if (shape_ == 1) return head[u] != head[v] ? order[0] : in[u] < in[v] ? u : v;
if (shape_ == 2) return u == v ? u : u == order[0] || v == order[0] ? order[0] : center_;
while (nodes[u].head != nodes[v].head) {
if (nodes[u].position > nodes[v].position) u = nodes[u].parent;
else v = nodes[v].parent;
}
return nodes[u].position < nodes[v].position ? u : v;
}
int dist(int u, int v) const {
int w = lca(u, v);
return depth[u] + depth[v] - 2 * depth[w];
}
bool is_ancestor(int u, int v) const {
check(u), check(v);
return in[u] <= in[v] && in[v] < out[u];
}
int ancestor(int v, int k) const {
check(v);
if (shape_) {
if (k < 0 || k > depth[v]) return -1;
if (!k) return v;
if (k == depth[v]) return order[0];
return shape_ == 2 ? parent[v] : order[in[v] - k];
}
if (k < 0) return -1;
int target = nodes[v].position + nodes[v].offset - k;
if (target < 0) return -1;
while (nodes[v].head + nodes[v].offset > target) v = nodes[v].parent;
return order[target - nodes[v].offset];
}
int jump(int u, int v, int k) const {
check(u), check(v);
if (k < 0) return -1;
if (!k) return u;
if (shape_ == 1) {
if (head[u] == head[v]) {
int delta = in[v] - in[u];
if (k > (delta < 0 ? -delta : delta)) return -1;
return order[in[u] + (delta < 0 ? -k : k)];
}
int a = depth[u], b = depth[v];
if (k > a + b) return -1;
if (k == a) return order[0];
return k < a ? order[in[u] - k] : order[in[v] - (a + b - k)];
}
if (shape_ == 2) {
int d = (u != center_) + (v != center_);
if (u == v || k > d) return -1;
return k == d ? v : center_;
}
int from = u, to = v, target = nodes[u].position + nodes[u].offset - k;
while (nodes[u].head != nodes[v].head) {
if (nodes[u].position > nodes[v].position) {
if (nodes[u].head + nodes[u].offset <= target) return order[target - nodes[u].offset];
u = nodes[u].parent;
} else v = nodes[v].parent;
}
int w = nodes[u].position < nodes[v].position ? u : v;
int d = nodes[w].position + nodes[w].offset;
int length = nodes[from].position + nodes[from].offset + nodes[to].position + nodes[to].offset - 2 * d;
if (k > length) return -1;
if (target >= d) return order[target - nodes[u].offset];
return ancestor(to, length - k);
}
std::pair<int, int> subtree(int v, bool edge = false) const {
check(v);
return {in[v] + edge, out[v]};
}
int edge_vertex(int id) const {
assert(0 <= id && id < int(edges.size()));
auto [u, v] = edges[id];
check(u);
return parent[u] == v ? u : v;
}
template <class F>
void for_each_path(int u, int v, F &&f, bool edge = false) const {
check(u), check(v);
if (shape_ == 1) {
int x = in[u], y = in[v];
if (head[u] == head[v]) {
int l = std::min(x, y) + edge, r = std::max(x, y) + 1;
if (l < r) f(l, r, x > y);
} else if (head[u] == order[0]) {
if (edge < x + 1) f(int(edge), x + 1, true);
f(in[head[v]], y + 1, false);
} else {
f(in[head[u]], x + 1, true);
if (edge < y + 1) f(int(edge), y + 1, false);
}
return;
}
std::array<std::array<int, 2>, 32> right;
int count = 0;
while (nodes[u].head != nodes[v].head) {
if (nodes[u].position > nodes[v].position) {
f(nodes[u].head, nodes[u].position + 1, true);
u = nodes[u].parent;
} else {
assert(count < int(right.size()));
right[count++] = {nodes[v].head, nodes[v].position + 1};
v = nodes[v].parent;
}
}
bool reversed = nodes[u].position > nodes[v].position;
int l = std::min(nodes[u].position, nodes[v].position) + edge, r = std::max(nodes[u].position, nodes[v].position) + 1;
if (l < r) f(l, r, reversed);
while (count) {
auto [a, b] = right[--count];
f(a, b, false);
}
}
std::vector<Segment> path(int u, int v, bool edge = false) const {
std::vector<Segment> result;
for_each_path(u, v, [&](int l, int r, bool reversed) { result.push_back({l, r, reversed}); }, edge);
return result;
}
};