永続順序付き多重集合

data_structure/persistent_multiset.hpp · GitHub で見る

共通ページ

重複を数えながら、過去版ごとの挿入・削除・順位・k 番目・前後の値を求めます。通常の集合として使う場合は、contains(root,x) が偽のときだけ挿入します。

#include "lib/data_structure/persistent_multiset.hpp"
#include <cassert>

int main() {
	PersistentMultiset<int> set;
	auto v0 = set.initial_root();
	auto v1 = set.insert(v0, 5, 3);
	auto v2 = set.insert(v1, 2);
	auto v3 = set.erase(v2, 5, 2);
	assert(set.size(v2) == 4 && set.count(v3, 5) == 1);
	assert(set.kth(v2, 0) == 2 && set.kth(v2, 3) == 5);
	assert(set.order_of_key(v2, 5) == 1);
	assert(set.lower_bound(v3, 3) == 5);
}

各更新は新しい root_type を返し、引数の版は変わりません。initial_root() が最初の版です。ルートをコピーするだけなら時間・追加領域ともに \(O(1)\) で、任意の過去版から分岐できます。同じオブジェクトから得たルートだけを渡してください。版ごとのノード解放は行いません。オブジェクト自体のコピーは全ノードの複製になります。

操作	内容	計算量
PersistentMultiset<T,Compare>(compare = Compare{})	空集合。既定は std::less<T>	\(O(1)\)
build(a)	vector<T> の全要素から版を作る。初期版は引き続き空集合	\(O(n\log n)\)
size(root) / empty(root)	重複込みの要素数／空か	\(O(1)\)
insert(root,x,count = 1)	x を count 個追加した版	\(O(\log d)\)
erase(root,x,count = 1)	x を最大 count 個削除した版。不足分は無視	\(O(\log d)\)
erase_all(root,x)	x を全て削除した版	\(O(\log d)\)
count(root,x) / contains(root,x)	x の個数／存在するか	\(O(\log d)\)
kth(root,k)	比較順で 0 始まり k 番目。重複込み	\(O(\log d)\)
order_of_key(root,x)	比較順で x より前にある要素数	\(O(\log d)\)
lower_bound(root,x)	x 以上の最初の値。なければ nullopt	\(O(\log d)\)
upper_bound(root,x)	x より大きい最初の値。なければ nullopt	\(O(\log d)\)

d は異なる値の個数です。個数・順位は std::int64_t で、count >= 0、全体の個数は INT64_MAX 以下とします。kth は 0 <= k < size(root) が必要です。Compare は狭義弱順序とし、等価な値は同じ値として数えます。lower_bound・upper_bound の大小も比較器に従います。T はコピー・代入可能な型とし、デフォルト構築は不要です。

reserve(count) は全版を合わせたノード数を予約します。node_count() は作成済みノード数、memory_usage() は予約済みノード領域のバイト数を返します。ルートは再確保後も有効です。値のコピー・演算は定数時間とします。更新の時間には領域確保の償却計算量を含みます。

一括構築は入力要素数を \(n\) として \(O(n\log n)\) 時間、追加ノード数は異なる値の個数 \(d\) 個です。挿入・削除・検索・順位取得は \(O(\log d)\) 時間、更新ごとの追加領域は \(O(\log d)\) です。更新では個数に比例する繰り返しは行いません。空集合・異なる値が 1 個の場合は定数時間です。読み取りはノードを作りません。

検証コード

• test/persistent_multiset.test.cpp

コード

#ifndef PERSISTENT_MULTISET_HPP
#define PERSISTENT_MULTISET_HPP

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <optional>
#include <utility>
#include <vector>

template <class T, class Compare = std::less<T>>
class PersistentMultiset {
  public:
	using value_type = T;
	using root_type = std::uint32_t;
  private:
	struct Node { std::optional<T> value; root_type left, right; std::int64_t count, size; int height; };
	std::vector<Node> nodes{{std::nullopt, 0, 0, 0, 0, 0}};
	Compare comp;
	void check([[maybe_unused]] root_type root) const { assert(root < nodes.size()); }
	root_type make(T value, root_type left, root_type right, std::int64_t count) {
		assert(nodes.size() < std::numeric_limits<root_type>::max());
		Node node{std::move(value), left, right, count, count + size(left) + size(right), std::max(nodes[left].height, nodes[right].height) + 1};
		auto result = static_cast<root_type>(nodes.size());
		nodes.push_back(std::move(node));
		return result;
	}
	root_type rebuild(root_type root, root_type left, root_type right, std::int64_t count) {
		assert(nodes.size() < std::numeric_limits<root_type>::max());
		auto result = static_cast<root_type>(nodes.size());
		nodes.push_back(nodes[root]);
		auto &node = nodes.back();
		node.left = left; node.right = right; node.count = count;
		node.size = count + nodes[left].size + nodes[right].size;
		node.height = std::max(nodes[left].height, nodes[right].height) + 1;
		return result;
	}
	root_type balance(root_type root, root_type left, root_type right, std::int64_t count) {
		if (nodes[left].height > nodes[right].height + 1) {
			auto a = nodes[left].left, b = nodes[left].right;
			if (nodes[a].height >= nodes[b].height) {
				auto r = rebuild(root, b, right, count);
				return rebuild(left, a, r, nodes[left].count);
			}
			auto l = rebuild(left, a, nodes[b].left, nodes[left].count);
			auto r = rebuild(root, nodes[b].right, right, count);
			return rebuild(b, l, r, nodes[b].count);
		}
		if (nodes[right].height > nodes[left].height + 1) {
			auto a = nodes[right].left, b = nodes[right].right;
			if (nodes[b].height >= nodes[a].height) {
				auto l = rebuild(root, left, a, count);
				return rebuild(right, l, b, nodes[right].count);
			}
			auto l = rebuild(root, left, nodes[a].left, count);
			auto r = rebuild(right, nodes[a].right, b, nodes[right].count);
			return rebuild(a, l, r, nodes[a].count);
		}
		return rebuild(root, left, right, count);
	}
	root_type erase_first(root_type root) {
		auto left = nodes[root].left, right = nodes[root].right;
		if (!left) return right;
		left = erase_first(left);
		return balance(root, left, right, nodes[root].count);
	}
	root_type update(root_type root, const T &value, std::int64_t count, bool remove) {
		if (!root) return remove ? 0 : make(value, 0, 0, count);
		auto left = nodes[root].left, right = nodes[root].right;
		auto next_count = nodes[root].count;
		if (comp(value, *nodes[root].value)) left = update(left, value, count, remove);
		else if (comp(*nodes[root].value, value)) right = update(right, value, count, remove);
		else {
			if (!remove) next_count += count;
			else if (count < next_count) next_count -= count;
			else {
				if (!left || !right) return left ? left : right;
				auto successor = right;
				while (nodes[successor].left) successor = nodes[successor].left;
				right = erase_first(right);
				return balance(successor, left, right, nodes[successor].count);
			}
		}
		if (left == nodes[root].left && right == nodes[root].right && next_count == nodes[root].count) return root;
		return balance(root, left, right, next_count);
	}

  public:
	explicit PersistentMultiset(Compare compare = Compare{}) : comp(std::move(compare)) {}
	root_type initial_root() const { return 0; }
	root_type build(std::vector<T> values) {
		assert(values.size() <= static_cast<std::size_t>(std::numeric_limits<std::int64_t>::max()));
		std::sort(values.begin(), values.end(), comp);
		std::vector<std::pair<T, std::int64_t>> distinct;
		for (const T &value : values) {
			if (distinct.empty() || comp(distinct.back().first, value)) distinct.emplace_back(value, 1);
			else ++distinct.back().second;
		}
		auto visit = [&](auto &&self, std::size_t l, std::size_t r) -> root_type {
			if (l == r) return 0;
			auto m = l + (r - l) / 2;
			auto left = self(self, l, m), right = self(self, m + 1, r);
			return make(distinct[m].first, left, right, distinct[m].second);
		};
		return visit(visit, 0, distinct.size());
	}
	std::int64_t size(root_type root) const { check(root); return nodes[root].size; }
	bool empty(root_type root) const { return size(root) == 0; }
	std::int64_t count(root_type root, const T &value) const {
		check(root);
		while (root) {
			const auto &node = nodes[root];
			if (comp(value, *node.value)) root = node.left;
			else if (comp(*node.value, value)) root = node.right;
			else return node.count;
		}
		return 0;
	}
	bool contains(root_type root, const T &value) const { return count(root, value) != 0; }
	root_type insert(root_type root, const T &value, std::int64_t count = 1) {
		assert(count >= 0 && size(root) <= std::numeric_limits<std::int64_t>::max() - count);
		return count ? update(root, value, count, false) : root;
	}
	root_type erase(root_type root, const T &value, std::int64_t count = 1) {
		check(root); assert(count >= 0);
		return count ? update(root, value, count, true) : root;
	}
	root_type erase_all(root_type root, const T &value) { return erase(root, value, std::numeric_limits<std::int64_t>::max()); }
	T kth(root_type root, std::int64_t k) const {
		assert(0 <= k && k < size(root));
		while (true) {
			const auto &node = nodes[root];
			auto left = size(node.left);
			if (k < left) root = node.left;
			else if (k < left + node.count) return *node.value;
			else { k -= left + node.count; root = node.right; }
		}
	}
	std::int64_t order_of_key(root_type root, const T &value) const {
		check(root);
		std::int64_t result = 0;
		while (root) {
			const auto &node = nodes[root];
			if (comp(*node.value, value)) { result += size(node.left) + node.count; root = node.right; }
			else root = node.left;
		}
		return result;
	}
	std::optional<T> lower_bound(root_type root, const T &value) const {
		check(root);
		std::optional<T> result;
		while (root) {
			const auto &node = nodes[root];
			if (!comp(*node.value, value)) { result = node.value; root = node.left; }
			else root = node.right;
		}
		return result;
	}
	std::optional<T> upper_bound(root_type root, const T &value) const {
		check(root);
		std::optional<T> result;
		while (root) {
			const auto &node = nodes[root];
			if (comp(value, *node.value)) { result = node.value; root = node.left; }
			else root = node.right;
		}
		return result;
	}
	void reserve(std::size_t count) { assert(count < std::numeric_limits<root_type>::max()); nodes.reserve(count + 1); }
	std::size_t node_count() const { return nodes.size() - 1; }
	std::size_t memory_usage() const { return nodes.capacity() * sizeof(Node); }
};

#endif

Back to top page