木上の Mo's Algorithm

tree/tree_mo.hpp · GitHub で見る

共通ページ

木の経路に含まれる頂点・辺を追加／削除して、経路上の異なる値の個数や頻度などを一括計算します。木と値は処理中に変更しません。集約は経路をたどる順番によらないものを扱います。

#include "lib/tree/tree_mo.hpp"
#include <vector>

int main() {
    std::vector<std::vector<int>> graph{ {1, 2}, {0, 3, 4}, {0}, {1}, {1} };
    std::vector<int> value{1, 2, 3, 4, 5}, answers(3);
    TreeMo mo(graph);
    mo.add_query(3, 2);
    mo.add_query(3, 4);
    mo.add_query(2, 2);
    int sum = 0;
    mo.run([&](int v) { sum += value[v]; },
        [&](int v) { sum -= value[v]; },
        [&](int id) { answers[id] = sum; });
    // answers = {10, 11, 3}
}

操作	内容・返り値
TreeMo(graph,root=0,edge_values=false)	std::vector<std::vector<int>> の無向木から構築
size() / query_count()	頂点数・問い合わせ数を int で返す
parent(v)	指定した根に対する親を int で返す。根は -1
add_query(u,v)	u と v を結ぶ経路を登録し、0 始まりの問い合わせ ID を int で返す
order(block_size=0)	処理順の問い合わせ ID を std::vector<int> で返す
run(add,erase,answer,block_size=0)	コールバックを呼んで処理。返り値は void

頂点番号は 0 始まりです。入力は自己ループ・多重辺のない連結な無向木とし、各辺を両方向の隣接リストに入れます。頂点数は INT_MAX/2 以下、空でない場合の root・u・v は [0,size()) にします。空の木も構築できますが、経路は登録できません。空の木では root を使いません。問い合わせ数は INT_MAX 以下です。

edge_values=false では両端を含む経路上の頂点を扱います。u==v の経路は 1 頂点です。edge_values=true では経路上の辺を扱い、コールバックにはその辺の子側の頂点番号を渡します。u==v の経路は空です。

辺値の場合は、上の例の mo と同じ木に対して次のように使えます。

TreeMo edges(graph, 0, true);
std::vector<int> weight{0, 10, 20, 30, 40};
// weight[v] は parent(v) と v を結ぶ辺の値。根の値は使いません。
edges.add_query(3, 2);
int edge_sum = 0, answer = 0;
edges.run([&](int v) { edge_sum += weight[v]; },
    [&](int v) { edge_sum -= weight[v]; },
    [&](int) { answer = edge_sum; });
// answer = 60

run の前に集約を空の状態にしてください。add(v) は頂点または辺を追加し、erase(v) は削除します。answer(id) の時点で、集約がその経路と一致します。LCA の追加・除外を呼び出し側で行う必要はありません。

コールバックの追加・削除順は経路上の順番とは限りません。非可換な経路積には Path Product を使ってください。処理中に問い合わせを追加しないでください。再実行は可能ですが、集約を空へ初期化する必要があります。終了時の集約を空へ戻すことは保証しません。

block_size は非負とし、0 では 通常の Mo’s Algorithm と同様に自動選択します。問い合わせがなければコールバックを呼びません。

頂点数を \(N\)、問い合わせ数を \(Q\) として、構築は \(O(N\log N)\) 時間・空間、経路の登録は償却 \(O(1)\)、順序生成は \(O(Q)\) です。コールバックが定数時間なら、自動選択時の処理は \(O(N\sqrt Q+Q)\) 時間、追加の作業空間は \(O(N+Q)\) です。幅 \(B\) を指定する場合の処理時間は \(O(QB+N^2/B+Q)\) です。

依存ライブラリ

• Default Init Allocator
• 64 bit キーの安定基数ソート
• Static Range Min
• Mo's Algorithm
• Sparse Table
• LCA

コード

#ifndef TREE_MO_HPP
#define TREE_MO_HPP

#include "../data_structure/mo_order.hpp"
#include "lca.hpp"

class TreeMo {
	LCA lca_;
	std::vector<int> enter_, leave_, tour_, extra_;
	std::vector<std::pair<int, int>> ranges_;
	bool edge_values_;

public:
	explicit TreeMo(const std::vector<std::vector<int>> &graph, int root = 0, bool edge_values = false)
		: lca_(), edge_values_(edge_values) {
		assert(graph.size() <= std::size_t(std::numeric_limits<int>::max() / 2));
		int n = static_cast<int>(graph.size());
		enter_.resize(n), leave_.resize(n);
		if (!n) return;
		assert(0 <= root && root < n);
		std::size_t edges = 0;
		for (const auto &row : graph) edges += row.size();
		assert(edges == std::size_t(n - 1) * 2);
		std::vector<int> parents(n, -2), stack{root};
		std::vector<std::size_t> next(n);
		parents[root] = -1;
		tour_.reserve(std::size_t(n) * 2);
		tour_.push_back(root);
		while (!stack.empty()) {
			int v = stack.back();
			if (next[v] < graph[v].size()) {
				int to = graph[v][next[v]++];
				assert(0 <= to && to < n);
				if (to == parents[v]) continue;
				assert(parents[to] == -2);
				parents[to] = v;
				enter_[to] = static_cast<int>(tour_.size());
				tour_.push_back(to);
				stack.push_back(to);
			} else {
				leave_[v] = static_cast<int>(tour_.size());
				tour_.push_back(v);
				stack.pop_back();
			}
		}
		assert(tour_.size() == std::size_t(n) * 2);
		lca_.build_from_parents(std::move(parents));
	}
	int size() const { return lca_.n; }
	int query_count() const { return static_cast<int>(ranges_.size()); }
	int parent(int v) const {
		assert(0 <= v && v < size());
		return lca_.parent[v];
	}
	int add_query(int u, int v) {
		assert(0 <= u && u < size() && 0 <= v && v < size());
		assert(ranges_.size() < std::size_t(std::numeric_limits<int>::max()));
		if (enter_[u] > enter_[v]) std::swap(u, v);
		int w = lca_.lca(u, v);
		if (w == u) {
			ranges_.emplace_back(enter_[u] + int(edge_values_), enter_[v] + 1);
			extra_.push_back(-1);
		} else {
			ranges_.emplace_back(leave_[u], enter_[v] + 1);
			extra_.push_back(edge_values_ ? -1 : w);
		}
		return query_count() - 1;
	}
	std::vector<int> order(int block_size = 0) const { return mo_order(ranges_, block_size); }
	template<class Add, class Erase, class Answer>
	void run(Add add, Erase erase, Answer answer, int block_size = 0) const {
		if (!query_count()) { assert(block_size >= 0); return; }
		std::vector<unsigned char> active(size());
		auto toggle = [&](int v) {
			if (active[v]) erase(v);
			else add(v);
			active[v] ^= 1;
		};
		int l = 0, r = 0;
		for (int id : order(block_size)) {
			auto [a, b] = ranges_[id];
			while (a < l) toggle(tour_[--l]);
			while (r < b) toggle(tour_[r++]);
			while (l < a) toggle(tour_[l++]);
			while (b < r) toggle(tour_[--r]);
			if (extra_[id] >= 0) toggle(extra_[id]);
			answer(id);
			if (extra_[id] >= 0) toggle(extra_[id]);
		}
	}
};

#endif

Back to top page