cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:heavy_check_mark: HLD (Heavy-Light Decomposition)

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
Back to top page