haar_lib/ds/
link_cut_tree.rs

1//! Link-Cut Tree
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/dynamic_tree_vertex_add_path_sum>
5//! - <https://judge.yosupo.jp/problem/dynamic_tree_vertex_set_path_composite> (非可換なモノイド)
6
7use std::{mem::swap, ptr};
8
9use crate::algebra::traits::*;
10
11struct Manager<M: Monoid> {
12    monoid: M,
13    nodes: Vec<Node<M>>,
14}
15
16struct Node<M: Monoid> {
17    value: M::Element,
18    acc: M::Element,
19    rev_acc: Option<M::Element>,
20    size: usize,
21    rev: bool,
22    lc: *mut Self,
23    rc: *mut Self,
24    par: *mut Self,
25}
26
27impl<M: Monoid + Commutative> Manager<M> {
28    fn create_commutative(&mut self) {
29        self.nodes.push(Node {
30            value: self.monoid.id(),
31            acc: self.monoid.id(),
32            rev_acc: None,
33            size: 1,
34            rev: false,
35            lc: ptr::null_mut(),
36            rc: ptr::null_mut(),
37            par: ptr::null_mut(),
38        })
39    }
40}
41
42impl<M: Monoid> Manager<M> {
43    fn new(monoid: M) -> Self {
44        Self {
45            monoid,
46            nodes: vec![],
47        }
48    }
49
50    fn create(&mut self) {
51        self.nodes.push(Node {
52            value: self.monoid.id(),
53            acc: self.monoid.id(),
54            rev_acc: Some(self.monoid.id()),
55            size: 1,
56            rev: false,
57            lc: ptr::null_mut(),
58            rc: ptr::null_mut(),
59            par: ptr::null_mut(),
60        })
61    }
62
63    fn ptr(&self, k: usize) -> *mut Node<M> {
64        &self.nodes[k] as *const _ as *mut _
65    }
66}
67
68impl<M: Monoid> Manager<M>
69where
70    M::Element: Clone,
71{
72    fn is_root(&self, this: *mut Node<M>) -> bool {
73        assert!(!this.is_null());
74        let par = unsafe { (*this).par };
75        par.is_null() || unsafe { (*par).lc != this && (*par).rc != this }
76    }
77
78    fn update(&self, this: *mut Node<M>) {
79        assert!(!this.is_null());
80        let this = unsafe { &mut *this };
81        let monoid = &self.monoid;
82        this.size = 1;
83        this.acc = this.value.clone();
84        if let Some(rev_acc) = this.rev_acc.as_mut() {
85            *rev_acc = this.value.clone();
86        }
87
88        let left = this.lc;
89        if !left.is_null() {
90            let left = unsafe { &mut *left };
91            self.pushdown(left);
92            this.size += left.size;
93            this.acc = monoid.op(left.acc.clone(), this.acc.clone());
94
95            if let Some(rev_acc) = this.rev_acc.as_mut() {
96                *rev_acc = monoid.op(rev_acc.clone(), left.rev_acc.clone().unwrap());
97            }
98        }
99
100        let right = this.rc;
101        if !right.is_null() {
102            let right = unsafe { &mut *right };
103            self.pushdown(right);
104            this.size += right.size;
105            this.acc = monoid.op(this.acc.clone(), right.acc.clone());
106
107            if let Some(rev_acc) = this.rev_acc.as_mut() {
108                *rev_acc = monoid.op(right.rev_acc.clone().unwrap(), rev_acc.clone());
109            }
110        }
111    }
112
113    fn reverse(&self, this: *mut Node<M>) {
114        if !this.is_null() {
115            unsafe { (*this).rev ^= true };
116        }
117    }
118
119    fn pushdown(&self, this: *mut Node<M>) {
120        assert!(!this.is_null());
121        let this = unsafe { &mut *this };
122        if this.rev {
123            swap(&mut this.lc, &mut this.rc);
124            self.reverse(this.lc);
125            self.reverse(this.rc);
126            this.rev = false;
127
128            if let Some(rev_acc) = this.rev_acc.as_mut() {
129                swap(&mut this.acc, rev_acc);
130            }
131        }
132    }
133
134    fn left_of(&self, this: *mut Node<M>) -> Option<*mut Node<M>> {
135        (!this.is_null()).then(|| unsafe { (*this).lc })
136    }
137
138    fn right_of(&self, this: *mut Node<M>) -> Option<*mut Node<M>> {
139        (!this.is_null()).then(|| unsafe { (*this).rc })
140    }
141
142    fn par_of(&self, this: *mut Node<M>) -> Option<*mut Node<M>> {
143        (!this.is_null()).then(|| unsafe { (*this).par })
144    }
145
146    fn set_left(&self, this: *mut Node<M>, left: *mut Node<M>) {
147        assert!(!this.is_null());
148        unsafe { (*this).lc = left };
149        if !left.is_null() {
150            unsafe { (*left).par = this };
151        }
152    }
153
154    fn set_right(&self, this: *mut Node<M>, right: *mut Node<M>) {
155        assert!(!this.is_null());
156        unsafe { (*this).rc = right };
157        if !right.is_null() {
158            unsafe { (*right).par = this };
159        }
160    }
161
162    fn set_par(&self, this: *mut Node<M>, par: *mut Node<M>) {
163        assert!(!this.is_null());
164        unsafe { (*this).par = par };
165    }
166
167    fn rot(&self, this: *mut Node<M>, dir_right: bool) {
168        let p = self.par_of(this).unwrap();
169        let g = self.par_of(p).unwrap();
170
171        if dir_right {
172            let c = self.right_of(this).unwrap();
173            self.set_left(p, c);
174            self.set_right(this, p);
175        } else {
176            let c = self.left_of(this).unwrap();
177            self.set_right(p, c);
178            self.set_left(this, p);
179        }
180
181        self.update(p);
182        self.update(this);
183
184        self.set_par(this, g);
185        if !g.is_null() {
186            let g = unsafe { &mut *g };
187            if g.lc == p {
188                g.lc = this;
189            }
190            if g.rc == p {
191                g.rc = this;
192            }
193            self.update(g);
194        }
195    }
196
197    fn splay(&self, this: *mut Node<M>) {
198        assert!(!this.is_null());
199        while !self.is_root(this) {
200            let p = self.par_of(this).unwrap();
201            let gp = self.par_of(p).unwrap();
202
203            if self.is_root(p) {
204                self.pushdown(p);
205                self.pushdown(this);
206                self.rot(this, this == self.left_of(p).unwrap());
207            } else {
208                self.pushdown(gp);
209                self.pushdown(p);
210                self.pushdown(this);
211                let flag = this == self.left_of(p).unwrap();
212
213                if (this == self.left_of(p).unwrap()) == (p == self.left_of(gp).unwrap()) {
214                    self.rot(p, flag);
215                    self.rot(this, flag);
216                } else {
217                    self.rot(this, flag);
218                    self.rot(this, !flag);
219                }
220            }
221        }
222        self.pushdown(this);
223    }
224
225    fn expose(&self, u: *mut Node<M>) {
226        let mut rp = ptr::null_mut();
227        let mut p = u;
228
229        while !p.is_null() {
230            self.splay(p);
231            unsafe { (*p).rc = rp };
232            self.update(p);
233            rp = p;
234            p = self.par_of(p).unwrap();
235        }
236
237        self.splay(u);
238        assert!(self.right_of(u).unwrap().is_null());
239    }
240
241    fn root_of(&self, mut this: *mut Node<M>) -> *mut Node<M> {
242        assert!(!this.is_null());
243        loop {
244            let p = self.par_of(this).unwrap();
245            if p.is_null() {
246                return this;
247            }
248            this = p;
249        }
250    }
251
252    fn same_group(&self, u: *mut Node<M>, v: *mut Node<M>) -> bool {
253        self.root_of(u) == self.root_of(v)
254    }
255
256    fn evert(&self, u: *mut Node<M>) {
257        self.expose(u);
258        self.reverse(u);
259        self.pushdown(u);
260    }
261}
262
263/// Link-Cut Tree
264///
265/// 動的に木の辺を追加・削除可能
266pub struct LinkCutTree<M: Monoid> {
267    man: Manager<M>,
268}
269
270impl<M: Monoid + Commutative + Clone> LinkCutTree<M> {
271    /// 可換モノイドを持てる`LinkCutTree<M>`を生成する。
272    pub fn new_commutative(monoid: M, n: usize) -> Self {
273        let mut man = Manager::new(monoid);
274        for _ in 0..n {
275            man.create_commutative();
276        }
277        Self { man }
278    }
279}
280
281impl<M: Monoid + Clone> LinkCutTree<M> {
282    /// 非可換なモノイドを持てる`LinkCutTree<M>`を生成する。
283    pub fn new_non_commutative(monoid: M, n: usize) -> Self {
284        let mut man = Manager::new(monoid);
285        for _ in 0..n {
286            man.create();
287        }
288        Self { man }
289    }
290}
291
292impl<M: Monoid + Clone> LinkCutTree<M>
293where
294    M::Element: Clone,
295{
296    #[allow(missing_docs)]
297    pub fn expose(&mut self, k: usize) {
298        let k = self.man.ptr(k);
299        self.man.expose(k);
300    }
301
302    /// 頂点`i`と頂点`j`の間にある辺を削除する。
303    ///
304    /// # Panics
305    ///
306    /// 頂点`i`と頂点`j`が隣接していないときパニックする。
307    pub fn cut(&mut self, i: usize, j: usize) {
308        assert_ne!(i, j);
309
310        let u = self.man.ptr(i);
311        let v = self.man.ptr(j);
312
313        self.man.evert(u);
314        self.man.expose(v);
315
316        assert!(self.man.left_of(u).unwrap().is_null());
317        assert!(self.man.is_root(v));
318
319        assert!(
320            self.man.left_of(v).unwrap() == u && self.man.par_of(u).unwrap() == v,
321            "`cut`操作では頂点`{i}`と`{j}`は隣接して連結されねばならない。",
322        );
323
324        self.man.set_par(u, ptr::null_mut());
325        self.man.set_left(v, ptr::null_mut());
326        self.man.update(v);
327    }
328
329    /// 頂点`i`と頂点`j`の間に辺を張る。
330    ///
331    /// # Panics
332    ///
333    /// 頂点`i`と頂点`j`が同一の木に属するときパニックする。
334    pub fn link(&mut self, i: usize, j: usize) {
335        assert_ne!(i, j);
336        let u = self.man.ptr(i);
337        let v = self.man.ptr(j);
338
339        self.man.expose(u);
340        self.man.evert(v);
341        let v = self.man.root_of(v);
342
343        assert!(
344            !self.man.same_group(u, v),
345            "`link`操作では頂点`{i}`と`{j}`は同一の木に属してはならない。",
346        );
347
348        assert!(self.man.is_root(u));
349        self.man.set_right(u, v);
350        self.man.update(u);
351    }
352
353    #[allow(missing_docs)]
354    pub fn evert(&mut self, k: usize) {
355        let k = self.man.ptr(k);
356        self.man.evert(k);
357    }
358
359    /// 頂点`k`の値を`x`に変更する。
360    pub fn set(&mut self, k: usize, x: M::Element) {
361        let u = self.man.ptr(k);
362        self.man.expose(u);
363        self.man.nodes[k].value = x;
364        self.man.update(u);
365    }
366
367    /// 頂点`k`の値をモノイドの演算と値`x`で更新する。
368    pub fn update(&mut self, k: usize, x: M::Element) {
369        self.set(k, self.man.monoid.op(self.get(k), x));
370    }
371
372    /// 頂点`k`の値を返す。
373    pub fn get(&self, k: usize) -> M::Element {
374        self.man.nodes[k].value.clone()
375    }
376
377    /// 頂点`i`,`j`間のパス上でのモノイドの演算の結果を返す。
378    ///
379    /// # Panics
380    ///
381    /// 頂点`i`と`j`が同一の木に属していないときパニックする。
382    pub fn fold(&self, i: usize, j: usize) -> M::Element {
383        let u = self.man.ptr(i);
384        let v = self.man.ptr(j);
385
386        self.man.evert(u);
387        self.man.expose(v);
388
389        assert!(
390            self.man.same_group(u, v),
391            "頂点`{i}`と`{j}`は同一の木に属していなければならない。",
392        );
393
394        self.man.nodes[j].acc.clone()
395    }
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401
402    use crate::algebra::trivial::*;
403
404    #[test]
405    fn test() {
406        let n = 10;
407
408        let mut lct = LinkCutTree::new_commutative(Trivial, n);
409
410        //        lct.cut(0, 1); // Runtime error
411
412        lct.link(0, 1);
413        lct.link(1, 2);
414
415        //        lct.link(0, 2); // Runtime error
416    }
417}