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