Skip to content

Commit f321e85

Browse files
hitonanodehitonanode
authored
Add auxiliary tree (online update) (#326)
* Add auxiliary tree (online update) * auxiliary tree: add test
1 parent 8ecb537 commit f321e85

File tree

3 files changed

+335
-0
lines changed

3 files changed

+335
-0
lines changed

‎tree/auxiliary_tree.hpp‎

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
#pragma once
2+
#include <unordered_set>
3+
#include <vector>
4+
5+
#include "../data_structure/fast_set.hpp"
6+
#include "../sparse_table/rmq_sparse_table.hpp"
7+
8+
// Data structure maintaining "compressed graph" of subsets of the vertices of a tree
9+
// Known as "auxiliary tree" and "virtual tree"
10+
// https://noshi91.github.io/algorithm-encyclopedia/auxiliary-tree
11+
class auxiliary_tree {
12+
13+
int n_ = 0;
14+
int root_ = -1;
15+
16+
// Each node is labeled by both v (given as input) and t (DFS preorder)
17+
std::vector<int> v2t; // v2t[v] = t
18+
std::vector<int> t2v; // t2v[t] = v
19+
20+
// To get LCA of two vertices in O(1) per query
21+
std::vector<int> _rmq_pos;
22+
StaticRMQ<std::pair<int, int>> _rmq;
23+
24+
// Auxiliary tree info
25+
// Variables starting with '_' are labeled by t, not v
26+
int _auxiliary_root = -1; // LCA of all currently activated vertices
27+
fast_set _is_active; // "t in _is_active" iff t is activated
28+
fast_set _is_semiactive; // "t in _is_semiactive" iff t is used in the current tree
29+
std::vector<int> _parent; // _parent[t] = parent of t in the current tree
30+
std::vector<std::unordered_set<int>> _child; // _child[t] = children of t in the current tree
31+
32+
int _get_lca(int t1, int t2) const {
33+
if (t1 > t2) std::swap(t1, t2);
34+
return _rmq.get(_rmq_pos.at(t1), _rmq_pos.at(t2) + 1).second;
35+
}
36+
37+
void _add_edge(int tpar, int tchild) {
38+
assert(tpar != tchild);
39+
assert(_parent.at(tchild) == -1);
40+
41+
_parent.at(tchild) = tpar;
42+
_child.at(tpar).insert(tchild);
43+
}
44+
45+
void _erase_edge(int tpar, int tchild) {
46+
assert(tpar != tchild);
47+
assert(_parent.at(tchild) == tpar);
48+
49+
_parent.at(tchild) = -1;
50+
_child.at(tpar).erase(tchild);
51+
}
52+
53+
public:
54+
int n() const { return n_; }
55+
56+
int original_root() const { return root_; }
57+
58+
int auxiliary_root() const { return _auxiliary_root == -1 ? -1 : t2v.at(_auxiliary_root); }
59+
60+
bool is_active(int v) const { return _is_active.contains(v2t.at(v)); }
61+
62+
bool is_semiactive(int v) const { return _is_semiactive.contains(v2t.at(v)); }
63+
64+
int get_parent(int v) const {
65+
const int t = v2t.at(v);
66+
return _parent.at(t) == -1 ? -1 : t2v.at(_parent.at(t));
67+
}
68+
69+
std::vector<int> get_children(int v) const {
70+
const int t = v2t.at(v);
71+
std::vector<int> ret;
72+
ret.reserve(_child.at(t).size());
73+
for (int c : _child.at(t)) ret.push_back(t2v.at(c));
74+
return ret;
75+
}
76+
77+
auxiliary_tree() = default;
78+
79+
auxiliary_tree(const std::vector<std::vector<int>> &to, int root)
80+
: n_(to.size()), root_(root), v2t(n_, -1), _rmq_pos(n_, -1), _is_active(n_),
81+
_is_semiactive(n_), _parent(n_, -1), _child(n_) {
82+
std::vector<std::pair<int, int>> dfspath; // (depth, t[v])
83+
t2v.reserve(n_);
84+
85+
auto rec = [&](auto &&self, int now, int prv, int depth) -> void {
86+
const int t = t2v.size();
87+
v2t.at(now) = t;
88+
t2v.push_back(now);
89+
90+
_rmq_pos.at(t) = dfspath.size();
91+
dfspath.emplace_back(depth, t);
92+
93+
for (int nxt : to.at(now)) {
94+
if (nxt == prv) continue;
95+
self(self, nxt, now, depth + 1);
96+
dfspath.emplace_back(depth, t);
97+
}
98+
};
99+
rec(rec, root, -1, 0);
100+
101+
_rmq = {dfspath, std::make_pair(n_, -1)};
102+
}
103+
104+
void activate(int v_) {
105+
const int t = v2t.at(v_);
106+
107+
if (_is_semiactive.contains(t)) {
108+
109+
// Already semiactive. Nothing to do.
110+
111+
} else if (_auxiliary_root == -1) {
112+
113+
// Add one vertex to empty set.
114+
_auxiliary_root = t;
115+
116+
} else if (const int next_root = _get_lca(_auxiliary_root, t); next_root != _auxiliary_root) {
117+
118+
// New node is outside the current tree. Update root.
119+
if (next_root != t) {
120+
_is_semiactive.insert(next_root);
121+
_add_edge(next_root, t);
122+
}
123+
_add_edge(next_root, _auxiliary_root);
124+
_auxiliary_root = next_root;
125+
126+
} else if (const int tnxt = _is_semiactive.next(t, n_);
127+
tnxt < n_ and _get_lca(t, tnxt) == t) {
128+
129+
// New node lies on the path of the current tree. Insert new node.
130+
const int tpar = _parent.at(tnxt);
131+
assert(tpar >= 0);
132+
133+
// tpar->tnxt => tpar->t->tnxt
134+
_erase_edge(tpar, tnxt);
135+
_add_edge(tpar, t);
136+
_add_edge(t, tnxt);
137+
138+
} else {
139+
140+
// New node is "under" the current tree.
141+
const int tprv = _is_semiactive.prev(t, -1);
142+
assert(tprv >= 0);
143+
const int tprvlca = _get_lca(t, tprv), tnxtlca = tnxt < n_ ? _get_lca(t, tnxt) : n_;
144+
145+
const int t2 = (tnxt == n_ or _get_lca(tprvlca, tnxtlca) == tnxtlca) ? tprv : tnxt;
146+
const int tlca = _get_lca(t, t2);
147+
148+
if (!_is_semiactive.contains(tlca)) {
149+
const int tc = _is_semiactive.next(tlca, n_);
150+
const int tpar = _parent.at(tc);
151+
assert(tpar >= 0);
152+
// tpar->tc => tpar->tlca->tc
153+
_is_semiactive.insert(tlca);
154+
_erase_edge(tpar, tc);
155+
_add_edge(tpar, tlca);
156+
_add_edge(tlca, tc);
157+
}
158+
159+
_add_edge(tlca, t);
160+
}
161+
162+
_is_semiactive.insert(t);
163+
_is_active.insert(t);
164+
}
165+
166+
void deactivate(int v_) {
167+
const int t = v2t.at(v_);
168+
169+
if (!_is_active.contains(t)) return;
170+
171+
const int num_children = _child.at(t).size();
172+
173+
if (num_children > 1) { // (1)
174+
175+
// Nothing to do (just deactivate it). Still semiactivated.
176+
177+
} else if (num_children == 1) {
178+
179+
// Delete this vertex from the current tree.
180+
const int tchild = *_child.at(t).begin();
181+
182+
if (_parent.at(t) == -1) {
183+
184+
// Root changes.
185+
// t->tchild => tchild
186+
_auxiliary_root = tchild;
187+
_erase_edge(t, tchild);
188+
189+
} else {
190+
191+
// tpar->t->tchild => tpar->tchild
192+
const int tpar = _parent.at(t);
193+
_erase_edge(tpar, t);
194+
_erase_edge(t, tchild);
195+
_add_edge(tpar, tchild);
196+
}
197+
198+
_is_semiactive.erase(t);
199+
200+
} else if (num_children == 0 and _parent.at(t) == -1) {
201+
202+
// Erase the only vertex in the current tree.
203+
_auxiliary_root = -1;
204+
_is_semiactive.erase(t);
205+
206+
} else {
207+
208+
assert(num_children == 0 and _parent.at(t) != -1);
209+
210+
const int tpar = _parent.at(t);
211+
const int tparpar = _parent.at(tpar);
212+
213+
if (!_is_active.contains(tpar) and _child.at(tpar).size() == 2) {
214+
215+
// In only this case, parent of t is also erased.
216+
const int t2 =
217+
t ^ (*_child.at(tpar).begin()) ^ (*std::next(_child.at(tpar).begin()));
218+
if (tparpar == -1) {
219+
// t<-tpar->t2 => t2
220+
_auxiliary_root = t2;
221+
_is_semiactive.erase(tpar);
222+
_erase_edge(tpar, t2);
223+
} else {
224+
// tparpar->tpar->t2 => tparpar->t2
225+
_is_semiactive.erase(tpar);
226+
_erase_edge(tparpar, tpar);
227+
_erase_edge(tpar, t2);
228+
_add_edge(tparpar, t2);
229+
}
230+
}
231+
_erase_edge(tpar, t);
232+
_is_semiactive.erase(t);
233+
}
234+
235+
_is_active.erase(t);
236+
}
237+
};

‎tree/auxiliary_tree.md‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
---
2+
title: LCA-based auxiliary tree / virtual tree, online （"虚树"）
3+
documentation_of: ./auxiliary_tree.hpp
4+
---
5+
6+
予め根付き木 $T$ が与えられる．$T$ の頂点部分集合 $S$ を $\emptyset$ で初期化した上で，以下のクエリをサポートする．
7+
8+
- $S$ に 1 頂点追加
9+
- $S$ から 1 頂点削除
10+
- $S$ の要素の組の最小共通祖先 (lowest common ancestor, LCA) 全てを頂点とし、もとの木と子孫関係を保った根付き木 $T'$ を考える． $T'$ に関して以下に答える：
11+
- $T$ の頂点 $v$ が $S$ に含まれるかどうか
12+
- $T$ の頂点 $v$ が $T'$ に含まれるかどうか
13+
- 特に $T'$ における $v$ の親
14+
- 特に $T'$ における $v$ の子の集合
15+
- $T'$ の根となる頂点
16+
17+
現実装では $T$ 上の LCA の計算を sparse table で行っているため， $\Theta(n \log n)$ の空間計算量を要する．
18+
19+
## 使用方法
20+
21+
```cpp
22+
vector<vector<int>> to(N); // edges of tree
23+
int root = 0;
24+
25+
auxiliary_tree at(to, root);
26+
27+
int v;
28+
at.activate(v); // Add v to S
29+
at.deactivate(v); // Remove v from S
30+
31+
int r = at.auxiliary_root(); // Root of T' (if exists) or -1
32+
33+
int par = at.get_parent(v); // Parent of v in T' (if exists) or -1
34+
vector<int> children = at.get_children(v); // Children of v in T'
35+
36+
bool b1 = at.is_active(v); // v in S?
37+
bool b2 = at.is_semiactive(v); // v in T'?
38+
```
39+
40+
## 問題例
41+
42+
- [鹿島建設プログラミングコンテスト2024（AtCoder Beginner Contest 340） G - Leaf Color](https://atcoder.jp/contests/abc340/tasks/abc340_g)
43+
- [No.901 K-ary εxtrεεmε - yukicoder](https://yukicoder.me/problems/no/901)

‎tree/test/auxiliary_tree.yuki901.test.cpp‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#define PROBLEM "https://yukicoder.me/problems/no/901"
2+
#include "../auxiliary_tree.hpp"
3+
#include "../../graph/shortest_path.hpp"
4+
#include <cassert>
5+
#include <iostream>
6+
using namespace std;
7+
8+
int main() {
9+
cin.tie(nullptr), ios::sync_with_stdio(false);
10+
int N;
11+
cin >> N;
12+
vector<vector<int>> to(N);
13+
14+
shortest_path<long long> sp(N);
15+
16+
for (int e = 0; e < N - 1; ++e) {
17+
int u, v, w;
18+
cin >> u >> v >> w;
19+
to.at(u).push_back(v);
20+
to.at(v).push_back(u);
21+
sp.add_bi_edge(u, v, w);
22+
}
23+
24+
const int root = 0;
25+
26+
auxiliary_tree at(to, root);
27+
28+
sp.solve(root);
29+
30+
int Q;
31+
cin >> Q;
32+
while (Q--) {
33+
int k;
34+
cin >> k;
35+
vector<int> xs(k);
36+
for (auto &x : xs) cin >> x;
37+
38+
for (int x : xs) at.activate(x);
39+
40+
long long ret = 0;
41+
42+
auto rec = [&](auto &&self, int now) -> void {
43+
for (int nxt : at.get_children(now)) {
44+
self(self, nxt);
45+
ret += sp.dist.at(nxt) - sp.dist.at(now);
46+
}
47+
};
48+
49+
rec(rec, at.auxiliary_root());
50+
51+
for (int x : xs) at.deactivate(x);
52+
53+
cout << ret << '\n';
54+
}
55+
}

0 commit comments

Comments
 (0)