Non-Recursive HLD Implementation

#	User	Rating
1	tourist	4009
2	jiangly	3831
3	Radewoosh	3646
4	jqdai0815	3620
4	Benq	3620
6	orzdevinwang	3529
7	ecnerwala	3446
8	Um_nik	3396
9	gamegame	3386
10	ksun48	3373

#	User	Contrib.
1	cry	164
1	maomao90	164
3	Um_nik	163
4	atcoder_official	160
5	-is-this-fft-	158
6	awoo	157
7	adamant	156
8	TheScrasse	154
8	nor	154
10	Dominater069	153

I wrote this pretty short implementation of Heavy-Light Decomposition today:

code

#include <iostream>
#include <vector>
using namespace std;

struct HLD {
	vector<int> par; // Parent
	vector<int> pp; // Path parent (Highest ancestor in same heavy-edge segment)
	vector<int> ind; // HLD index (Index in the array HLD maps nodes to)

	// p: parent of node i. Must have p[i] < i (p[0] = -1).
	HLD(const vector<int>& p) : par(p), pp(p.size()), ind(p.size(), -1) {
		int n = p.size();
		vector<int> siz(n, 1); // subtree size
		for (int i = n-1; i > 0; --i) siz[par[i]] += siz[i];

		vector<int> pc(n, -1); // Preferred child
		for (int i = n-1; i > 0; --i) {
			if (2*siz[i] >= siz[par[i]]) pc[par[i]] = i;
		}

		int cur = 0; // Current position in array
		for (int i = 0; i < n; ++i) {
			if (ind[i] != -1) continue;
			for (int j = i; j != -1; j = pc[j], ++cur) {
				ind[j] = cur;
				pp[j] = i;
			}
		}
	}
	// Get intervals corresponding to path between a and b
	vector<pair<int, int>> get(int a, int b) {
		vector<pair<int, int>> res;
		while(true) {
			if (ind[b] < ind[a]) swap(a, b);
			if (ind[pp[b]] <= ind[a]) {
				res.push_back({ind[a], ind[b]});
				return res;
			} else {
				res.push_back({ind[pp[b]], ind[b]});
				b = par[pp[b]];
			}
		}
	}
};

It uses no recursion, which is nice, and takes data in the index-of-parent format, which can be a small plus or a large minus (I needed HLD over Aho-Corasick automata, so that input format is useful there). The intervals returned by get are in decreasing order, not in the order they are in the path, so something like adding an arithmetic sequence to a path is not possible.

The par-array just stores the parent of every node, ind gives the index of the node in the array HLD maps the nodes to, and pp gives the highest ancestor of the node that is mapped to the same continuous segment.

get returns a vector of at most $$$O(log(n))$$$ intervals, such that hld-index of every node on the path between $$$a$$$ and $$$b$$$ is in exactly one of the returned intervals. Specifically it returns all intersections of heavy paths with the path between $$$a$$$ and $$$b$$$ that are nonempty, in decreasing order (first interval has highest startpoint and endpoint). Clearly there are at most $O(log(n))$ such intervals, since every path contains at most $O(log(n))$ light edges, so at most $O(log(n))$ heavy intervals.

Get works since the path from $$$b$$$ to $$$pp[b]$$$ corresponds to the continuous interval $$$ind[pp[b]], ind[b]$$$ in the array HLD maps the nodes to, and therefore if $$$ind[pp[b]] \leq ind[a] \leq ind[b]$$$, $$$a$$$ is on the path between $$$pp[b]$$$ and $$$b$$$, and if $$$ind[a] < ind[pp[b]]$$$, then no node we go over is an ancestor of $$$a$$$, since the HLD-index of a node's ancestor is always less than the ind of the node.

Since get also finds the LCA, this also gives a nice LCA implementation:

code

        // Get LCA of a and b
        int lca(int a, int b) {
                while(true) {
                        if (ind[b] < ind[a]) swap(a, b); 
                        if (ind[pp[b]] <= ind[a]) return a;
                        else b = par[pp[b]];
                }
        }

Here's Caves and Tunnels solved with this HLD-implementation. Changing the tree format is pretty nasty, so the code isn't that clean :/

code

#include <iostream>
#include <vector>
using namespace std;
using ll = long long;

struct HLD {
	vector<int> par; // Parent
	vector<int> pp; // Path parent
	vector<int> ind; // HLD index

	// p: parent of node i. Must have p[i] < i (p[0] = -1).
	HLD(const vector<int>& p) : par(p), pp(p.size()), ind(p.size(), -1) {
		int n = p.size();
		vector<int> siz(n, 1); // subtree size
		for (int i = n-1; i > 0; --i) siz[par[i]] += siz[i];

		vector<int> pc(n, -1); // Preferred child
		for (int i = n-1; i > 0; --i) {
			if (2*siz[i] >= siz[par[i]]) pc[par[i]] = i;
		}

		int cur = 0; // Current position in array
		for (int i = 0; i < n; ++i) {
			if (ind[i] != -1) continue;
			for (int j = i; j != -1; j = pc[j], ++cur) {
				ind[j] = cur;
				pp[j] = i;
			}
		}
	}
	// Get intervals corresponding to path between a and b
	vector<pair<int, int>> get(int a, int b) {
		vector<pair<int, int>> res;
		while(true) {
			if (ind[b] < ind[a]) swap(a, b);
			if (ind[pp[b]] <= ind[a]) {
				res.push_back({ind[a], ind[b]});
				return res;
			} else {
				res.push_back({ind[pp[b]], ind[b]});
				b = par[pp[b]];
			}
		}
	}
};

struct SegTree {
	vector<ll> seg;
	int h = 1;

	SegTree(int n) {
		while(h < n) h <<= 1;
		seg.resize(2*h, 0);
	}
	ll recGet(int a, int b, int i, int ia, int ib) const {
		if (b < ia || ib < a) return 0;
		if (a <=ia && ib <=b) return seg[i];
		int im = (ia + ib) >> 1;
		ll res = 0;
		res = max(res, recGet(a, b, 2*i, ia, im));
		res = max(res, recGet(a, b, 2*i+1, im+1, ib));
		return res;
	}
	ll get(int a, int b) const {
		return recGet(a, b, 1, 0, h-1);
	}
	void add(int i, int v) {
		i += h;
		seg[i] += v;
		for (i >>= 1; i >= 1; i >>= 1) {
			seg[i] = max(seg[2*i], seg[2*i+1]);
		}
	}
};

const int N = 1e5;
vector<int> conns[N];
int ord[N];

int dfs(int i, int p = -1, int j = 0) {
	ord[i] = j;
	j += 1;
	for (auto t : conns[i]) {
		if (t != p) j = dfs(t, i, j);
	}
	return j;
}

int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(0);

	int n;
	cin >> n;
	for (int i = 0; i < n-1; ++i) {
		int a, b;
		cin >> a >> b;
		--a; --b;
		conns[a].push_back(b);
		conns[b].push_back(a);
	}

	dfs(0);
	vector<int> par(n, -1);
	for (int i = 0; i < n; ++i) {
		for (auto t : conns[i]) {
			if (ord[t] < ord[i]) par[ord[i]] = ord[t];
		}
	}

	HLD hld(par);
	SegTree seg(n);

	int q;
	cin >> q;
	for (int j = 0; j < q; ++j) {
		char t;
		cin >> t;
		if (t == 'I') {
			int i, v;
			cin >> i >> v;
			i = ord[i-1];
			seg.add(hld.ind[i], v);
		} else {
			int a, b;
			cin >> a >> b;
			a = ord[a-1]; b = ord[b-1];

			auto ints = hld.get(a, b);
			ll res = 0;
			for (auto pr : ints) {
				res = max(res, seg.get(pr.first, pr.second));
			}
			cout << res << '\n';
		}
	}
}

	Rev.	Lang.	By	When	Δ	Comment
	en2		mango_lassi	2019-05-22 07:59:19	7	Tiny change: '\n ind[i] = cur;\n' -> '\n ind[j] = cur;\n'I had fixed a bug in the Caves and Tunnels version of HLD that still existed in the first code.
	en1		mango_lassi	2019-05-21 19:50:30	6388	Initial revision (published)

History