Priority Search Tree（三辺矩形内の点の列挙）

data_structure/priority_search_tree.hpp · GitHub で見る

共通ページ

二次元の点について、left <= x < right かつ y < upper を満たす点を列挙します。x 区間内で y が最小の点の検索、点の挿入・削除、y 座標の変更にも対応します。

x 座標は構築時に登録し、各点を入力順の 0 始まりの ID で指定します。同じ x 座標や同じ座標の点も、異なる ID として登録できます。構築後に新しい x 座標や ID を追加する場合は build または build_coordinates で作り直します。

#include "lib/data_structure/priority_search_tree.hpp"
#include <iostream>

std::vector<std::pair<int, int>> points = { {1, 5}, {3, 2}, {3, 7}, {6, 1} };
PrioritySearchTree<int> tree(points);
auto ids = tree.report(1, 6, 6); // ID 0, 1（順序は未規定）
auto best = tree.min_y(1, 6);    // optional<int>(1)
tree.set(2, 0);                 // ID 2 の点を (3, 0) に変更
tree.erase(1);
tree.enumerate(0, 10, 3, [&](int id) {
	std::cout << id << ' ' << tree.x(id) << ' ' << *tree.get(id) << '\n';
});

PrioritySearchTree<int> active(std::vector<int>{2, 4, 6}); // x 座標だけ登録
active.insert(1, 8); // ID 1 の点 (4, 8) を挿入

登録した点数を \(n\)、列挙する点数を \(k\) とします。

操作	内容・返り値	計算量
PrioritySearchTree<X=int64_t,Y=X>()	登録点のない構造を作る	\(O(1)\)
PrioritySearchTree(points) / build(points)	vector<pair<X,Y>> を登録し、全点を有効にする	\(O(n\log n)\)
PrioritySearchTree(coordinates) / build_coordinates(coordinates)	vector<X> を登録し、全点を無効にする	\(O(n\log n)\)
point_count() / size() / empty()	登録点数 / 有効な点数 / 有効な点がないか	\(O(1)\)
x(id) / get(id) / contains(id)	x 座標 / optional<Y> / 有効かどうか	\(O(1)\)
insert(id,y)	無効な ID に点を挿入し true を返す。既に有効なら変更せず false	\(O(\log n)\)
erase(id)	点を無効にする。削除できたかを返す	\(O(\log n)\)
set(id,y)	y 座標を変更する。無効なら挿入する	\(O(\log n)\)
min_y(left,right) const	[left,right) 内で y が最小の有効な点の ID。存在しなければ nullopt	\(O(\log n)\)
report(left,right,upper) const	条件を満たす有効な点の ID を vector<int> で返す	\(O(\log n+k)\)
enumerate(left,right,upper,visit) const	条件を満たす各 ID に visit(id) を呼ぶ	\(O(\log n+k)\)
clear()	登録した x 座標と ID を保ち、全点を無効にする	\(O(n)\)
swap(other)	内容を交換する	\(O(1)\)

min_y で y が等しい場合は最小の ID を返します。列挙の順序は未規定で、同じ ID は一度だけ列挙します。left >= right の問い合わせは空です。enumerate のコールバック内で対象の木を変更しないでください。

ID の引数は 0 <= id < point_count() とします。x と get は読み取り専用の参照を返します。ただし X=bool の x は値を返します。X・Y にはコピーと < による厳密弱順序が必要で、デフォルト構築や無限大の値は不要です。浮動小数点型では NaN を渡さないでください。登録点数は INT_MAX 以下とします。

保持空間は \(O(n)\) です。enumerate・min_y の追加空間は \(O(\log n)\)、report は返り値を含め \(O(\log n+k)\) です。比較・コピーとコールバックを定数時間とします。空の構造への問い合わせも利用できます。

参考: McCreight, Priority Search Trees。

コード

#ifndef PRIORITY_SEARCH_TREE_HPP
#define PRIORITY_SEARCH_TREE_HPP

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

template <class X = std::int64_t, class Y = X>
class PrioritySearchTree {
	std::vector<X> xs;
	std::vector<std::optional<Y>> ys;
	std::vector<int> order, ranks, heap;
	std::vector<std::size_t> positions;
	std::size_t base = 1;
	int active = 0;
	bool less(int a, int b) const {
		if (a < 0) return false;
		if (b < 0) return true;
		if (*ys[a] < *ys[b]) return true;
		if (*ys[b] < *ys[a]) return false;
		return a < b;
	}
	void push(int id) {
		std::size_t v = 1, left = 0, right = base;
		while (heap[v] >= 0) {
			if (less(id, heap[v])) {
				std::swap(id, heap[v]);
				positions[heap[v]] = v;
			}
			std::size_t middle = left + (right - left) / 2;
			if (std::size_t(ranks[id]) < middle) v *= 2, right = middle;
			else v = v * 2 + 1, left = middle;
		}
		heap[v] = id;
		positions[id] = v;
	}
	void fill_hole(std::size_t v) {
		while (v < base) {
			std::size_t child = v * 2;
			if (less(heap[child + 1], heap[child])) ++child;
			heap[v] = heap[child];
			if (heap[v] < 0) return;
			positions[heap[v]] = v;
			v = child;
		}
		heap[v] = -1;
	}
	std::pair<std::size_t, std::size_t> interval(const X &left, const X &right) const {
		if (!(left < right)) return {0, 0};
		auto lower = [&](const X &x) {
			return std::size_t(std::lower_bound(order.begin(), order.end(), x,
				[&](int id, const X &value) { return xs[id] < value; }) - order.begin());
		};
		return {lower(left), lower(right)};
	}
	template <class F>
	void enumerate_all(std::size_t v, const Y &upper, F &visit) const {
		int id = heap[v];
		if (id < 0 || !(*ys[id] < upper)) return;
		visit(id);
		if (v < base) {
			enumerate_all(v * 2, upper, visit);
			enumerate_all(v * 2 + 1, upper, visit);
		}
	}
	template <class F>
	void enumerate_node(std::size_t v, std::size_t left, std::size_t right,
		std::size_t ql, std::size_t qr, const Y &upper, F &visit) const {
		int id = heap[v];
		if (id < 0 || !(*ys[id] < upper)) return;
		if (ql <= left && right <= qr) {
			enumerate_all(v, upper, visit);
			return;
		}
		if (ql <= std::size_t(ranks[id]) && std::size_t(ranks[id]) < qr) visit(id);
		std::size_t middle = left + (right - left) / 2;
		if (ql < middle) enumerate_node(v * 2, left, middle, ql, qr, upper, visit);
		if (middle < qr) enumerate_node(v * 2 + 1, middle, right, ql, qr, upper, visit);
	}
	int min_node(std::size_t v, std::size_t left, std::size_t right,
		std::size_t ql, std::size_t qr) const {
		int id = heap[v];
		if (id < 0 || (ql <= left && right <= qr)) return id;
		int best = ql <= std::size_t(ranks[id]) && std::size_t(ranks[id]) < qr ? id : -1;
		std::size_t middle = left + (right - left) / 2;
		if (ql < middle) {
			int candidate = min_node(v * 2, left, middle, ql, qr);
			if (less(candidate, best)) best = candidate;
		}
		if (middle < qr) {
			int candidate = min_node(v * 2 + 1, middle, right, ql, qr);
			if (less(candidate, best)) best = candidate;
		}
		return best;
	}
	void check_id([[maybe_unused]] int id) const { assert(0 <= id && id < point_count()); }

public:
	PrioritySearchTree() = default;
	explicit PrioritySearchTree(const std::vector<std::pair<X, Y>> &points) { build(points); }
	explicit PrioritySearchTree(const std::vector<X> &coordinates) { build_coordinates(coordinates); }
	PrioritySearchTree(const PrioritySearchTree &) = default;
	PrioritySearchTree &operator=(const PrioritySearchTree &) = default;
	PrioritySearchTree(PrioritySearchTree &&other) noexcept : PrioritySearchTree() { swap(other); }
	PrioritySearchTree &operator=(PrioritySearchTree &&other) noexcept {
		if (this != &other) {
			PrioritySearchTree replacement(std::move(other));
			swap(replacement);
		}
		return *this;
	}
	void swap(PrioritySearchTree &other) noexcept {
		xs.swap(other.xs);
		ys.swap(other.ys);
		order.swap(other.order);
		ranks.swap(other.ranks);
		heap.swap(other.heap);
		positions.swap(other.positions);
		std::swap(base, other.base);
		std::swap(active, other.active);
	}
	int size() const { return active; }
	bool empty() const { return active == 0; }
	int point_count() const { return int(xs.size()); }
	typename std::vector<X>::const_reference x(int id) const { check_id(id); return xs[id]; }
	const std::optional<Y> &get(int id) const { check_id(id); return ys[id]; }
	bool contains(int id) const { check_id(id); return positions[id] != 0; }
	void build_coordinates(const std::vector<X> &coordinates) {
		assert(coordinates.size() <= std::size_t(std::numeric_limits<int>::max()));
		xs = coordinates;
		ys.clear();
		ys.resize(xs.size());
		positions.assign(xs.size(), 0);
		order.resize(xs.size());
		ranks.resize(xs.size());
		std::iota(order.begin(), order.end(), 0);
		std::sort(order.begin(), order.end(), [&](int a, int b) {
			if (xs[a] < xs[b]) return true;
			if (xs[b] < xs[a]) return false;
			return a < b;
		});
		for (int i = 0; i < point_count(); ++i) ranks[order[i]] = i;
		base = 1;
		while (base < xs.size()) base *= 2;
		heap.assign(base * 2, -1);
		active = 0;
	}
	void build(const std::vector<std::pair<X, Y>> &points) {
		std::vector<X> coordinates;
		coordinates.reserve(points.size());
		for (const auto &point : points) coordinates.push_back(point.first);
		build_coordinates(coordinates);
		for (int id = 0; id < point_count(); ++id) {
			ys[id] = points[id].second;
			heap[base + ranks[id]] = id;
			positions[id] = base + ranks[id];
		}
		for (std::size_t v = base; --v;) fill_hole(v);
		active = point_count();
	}
	void clear() {
		std::fill(heap.begin(), heap.end(), -1);
		std::fill(positions.begin(), positions.end(), 0);
		for (auto &y : ys) y.reset();
		active = 0;
	}
	bool insert(int id, Y y) {
		check_id(id);
		if (contains(id)) return false;
		ys[id] = std::move(y);
		push(id);
		++active;
		return true;
	}
	bool erase(int id) {
		check_id(id);
		if (!contains(id)) return false;
		fill_hole(positions[id]);
		positions[id] = 0;
		ys[id].reset();
		--active;
		return true;
	}
	void set(int id, Y y) {
		erase(id);
		insert(id, std::move(y));
	}
	std::optional<int> min_y(const X &left, const X &right) const {
		auto [l, r] = interval(left, right);
		if (l == r) return std::nullopt;
		int id = min_node(1, 0, base, l, r);
		return id < 0 ? std::nullopt : std::optional<int>(id);
	}
	template <class F>
	void enumerate(const X &left, const X &right, const Y &upper, F &&visit) const {
		auto [l, r] = interval(left, right);
		if (l < r) enumerate_node(1, 0, base, l, r, upper, visit);
	}
	std::vector<int> report(const X &left, const X &right, const Y &upper) const {
		std::vector<int> result;
		enumerate(left, right, upper, [&](int id) { result.push_back(id); });
		return result;
	}
};

#endif

Back to top page