library/data-structure/persistent-dsu.hpp

• View this file on GitHub
• Last update: 2023-08-13 14:16:40+08:00
• Include: #include "library/data-structure/persistent-dsu.hpp"

Depends on

• library/data-structure/persistent-array.hpp

Verified with

• test/data-structure/persistent-dsu/yosupo-Persistent-Unionfind.test.cpp
• test/data-structure/persistent-dsu/yosupo-Unionfind.test.cpp

Code

#pragma once
#include <vector>

#include <cassert>

#include "persistent-array.hpp"


namespace felix {

struct persistent_dsu {
public:
	persistent_dsu() : n(0) {}
	explicit persistent_dsu(int _n) : n(_n), data(std::vector<int>(_n, -1)) {}

	int versions() const { return data.versions(); }

private:
	struct ref {
		int id;
		int n;
		persistent_array<int>& data;

		int leader(int u) const {
			int p;
			while((p = data[id].get(u)) >= 0) {
				u = p;
			}
			return u;
		}

		bool same(int u, int v) const {
			return leader(u) == leader(v);
		}

		int size(int u) const {
			int p;
			do {
				p = data[id].get(u);
				if(p < 0) {
					return -p;
				}
				u = p;
			} while(true);
		}

		std::pair<int, bool> merge(int u, int v) {
			u = leader(u), v = leader(v);
			if(u == v) {
				return {id, false};
			}
			int sz_u = size(u), sz_v = size(v);
			if(sz_u < sz_v) {
				std::swap(u, v);
				std::swap(sz_u, sz_v);
			}
			int new_id = data[id].set(v, u);
			new_id = data[new_id].set(u, -(sz_u + sz_v));
			return {new_id, true};
		}
	};

public:
	ref operator[](int id) {
		assert(0 <= id && id < data.versions());
		return ref{id, n, data};
	}

private:
	int n;
	persistent_array<int> data;
};

} // namespace felix

#line 2 "library/data-structure/persistent-dsu.hpp"
#include <vector>

#include <cassert>

#line 3 "library/data-structure/persistent-array.hpp"
#include <algorithm>

#include <functional>

#line 6 "library/data-structure/persistent-array.hpp"

namespace felix {

template<class T>
struct persistent_array {
public:
	persistent_array() : n(0) {}
	explicit persistent_array(int _n) : persistent_array(std::vector<T>(_n)) {}
	explicit persistent_array(const std::vector<T>& v) : n(v.size()) {
		std::function<node_t*(int, int)> build = [&](int L, int R) {
			if(L + 1 == R) {
				return new node_t(v[L]);
			}
			int mid = (L + R) / 2;
			return new node_t(build(L, mid), build(mid, R));
		};
		roots.push_back(build(0, n));
	}

	int versions() const { return (int) roots.size(); }

private:
	struct node_t {
		T val = T();

		node_t* lc = nullptr;
		node_t* rc = nullptr;

		node_t() {}
		node_t(const T& x) : val(x) {}
		node_t(node_t* ll, node_t* rr) : lc(ll), rc(rr) {}

		node_t* set(int pos, const T& x, int L, int R) {
			if(L + 1 == R) {
				return new node_t(x);
			}
			auto v = new node_t(*this);
			int mid = (L + R) / 2;
			if(pos < mid) {
				v->lc = v->lc->set(pos, x, L, mid);
			} else {
				v->rc = v->rc->set(pos, x, mid, R);
			}
			return v;
		}

		T get(int pos, int L, int R) {
			if(L + 1 == R) {
				return val;
			}
			int mid = (L + R) / 2;
			if(pos < mid) {
				return lc->get(pos, L, mid);
			} else {
				return rc->get(pos, mid, R);
			}
		}
	};

	int n;
	std::vector<node_t*> roots;

	struct tree_ref {
		node_t* root;
		int n;
		std::vector<node_t*>& roots;

		int set(int pos, const T& x) {
			assert(0 <= pos && pos < n);
			roots.push_back(root->set(pos, x, 0, n));
			return roots.size() - 1;
		}

		T get(int pos) const {
			assert(0 <= pos && pos < n);
			return root->get(pos, 0, n);
		}

		T operator[](int pos) const { return get(pos); }
	};

public:
	tree_ref operator[](int id) {
		assert(0 <= id && id < (int) roots.size());
		return tree_ref{roots[id], n, roots};
	}
};

} // namespace felix

#line 5 "library/data-structure/persistent-dsu.hpp"

namespace felix {

struct persistent_dsu {
public:
	persistent_dsu() : n(0) {}
	explicit persistent_dsu(int _n) : n(_n), data(std::vector<int>(_n, -1)) {}

	int versions() const { return data.versions(); }

private:
	struct ref {
		int id;
		int n;
		persistent_array<int>& data;

		int leader(int u) const {
			int p;
			while((p = data[id].get(u)) >= 0) {
				u = p;
			}
			return u;
		}

		bool same(int u, int v) const {
			return leader(u) == leader(v);
		}

		int size(int u) const {
			int p;
			do {
				p = data[id].get(u);
				if(p < 0) {
					return -p;
				}
				u = p;
			} while(true);
		}

		std::pair<int, bool> merge(int u, int v) {
			u = leader(u), v = leader(v);
			if(u == v) {
				return {id, false};
			}
			int sz_u = size(u), sz_v = size(v);
			if(sz_u < sz_v) {
				std::swap(u, v);
				std::swap(sz_u, sz_v);
			}
			int new_id = data[id].set(v, u);
			new_id = data[new_id].set(u, -(sz_u + sz_v));
			return {new_id, true};
		}
	};

public:
	ref operator[](int id) {
		assert(0 <= id && id < data.versions());
		return ref{id, n, data};
	}

private:
	int n;
	persistent_array<int> data;
};

} // namespace felix

Back to top page