competitive/data_structure/
splay_operations.rs

1use super::binary_search_tree::{
2    BstDataMutRef, BstNodePtr, BstRoot, BstSeeker, BstSpec, EqualSide,
3    node::{WithNoParent, WithParent},
4    seeker::SeekRight,
5};
6use std::{cmp::Ordering, mem::MaybeUninit};
7
8pub mod with_parent {
9    use super::{BstDataMutRef, BstNodePtr, BstSpec, MaybeUninit, WithParent};
10
11    type NodePtr<Spec> = BstNodePtr<<Spec as BstSpec>::Data, <Spec as BstSpec>::Parent>;
12
13    #[inline]
14    unsafe fn internal_parent<Spec, Data>(
15        node: NodePtr<Spec>,
16    ) -> Result<(NodePtr<Spec>, usize), Option<NodePtr<Spec>>>
17    where
18        Spec: BstSpec<Data = Data, Parent = WithParent<Data>>,
19    {
20        let Some(parent) = (unsafe { node.as_ref().parent.parent }) else {
21            return Err(None);
22        };
23        let children = unsafe { parent.as_ref().child };
24        if children[0] == Some(node) {
25            Ok((parent, 0))
26        } else if children[1] == Some(node) {
27            Ok((parent, 1))
28        } else {
29            Err(Some(parent))
30        }
31    }
32
33    #[inline(always)]
34    unsafe fn rotate_at<Spec, Data>(
35        mut node: NodePtr<Spec>,
36        mut parent: NodePtr<Spec>,
37        direction: usize,
38    ) where
39        Spec: BstSpec<Data = Data, Parent = WithParent<Data>>,
40    {
41        let middle = unsafe { node.as_ref().child[direction ^ 1] };
42        unsafe {
43            parent.as_mut().child[direction] = middle;
44            if let Some(mut middle) = middle {
45                middle.as_mut().parent.parent = Some(parent);
46            }
47            node.as_mut().child[direction ^ 1] = Some(parent);
48            parent.as_mut().parent.parent = Some(node);
49        }
50    }
51
52    #[inline(always)]
53    pub unsafe fn rotate<Spec, Data>(mut node: NodePtr<Spec>)
54    where
55        Spec: BstSpec<Data = Data, Parent = WithParent<Data>>,
56    {
57        let (parent, direction) = unsafe { internal_parent::<Spec, Data>(node) }
58            .expect("an auxiliary root cannot be rotated");
59        unsafe {
60            node.as_mut().parent.parent = parent.as_ref().parent.parent;
61            if let Ok((mut grandparent, direction)) = internal_parent::<Spec, Data>(parent) {
62                grandparent.as_mut().child[direction] = Some(node);
63            }
64            rotate_at::<Spec, Data>(node, parent, direction);
65            Spec::bottom_up(BstDataMutRef::new_unchecked(parent));
66        }
67    }
68
69    /// Moves `node` to the root of its auxiliary tree and returns the previous root.
70    ///
71    /// # Safety
72    ///
73    /// `node` and every pointer reachable through its auxiliary-parent chain must
74    /// refer to live nodes of the same tree.
75    #[inline(always)]
76    pub unsafe fn splay<Spec, Data>(node: NodePtr<Spec>) -> NodePtr<Spec>
77    where
78        Spec: BstSpec<Data = Data, Parent = WithParent<Data>>,
79    {
80        let mut inline_stack = [const { MaybeUninit::uninit() }; 64];
81        let mut inline_len = 0;
82        let mut overflow_stack = Vec::new();
83        let mut current = node;
84        loop {
85            if inline_len < inline_stack.len() {
86                inline_stack[inline_len].write(current);
87                inline_len += 1;
88            } else {
89                overflow_stack.push(current);
90            }
91            match unsafe { internal_parent::<Spec, Data>(current) } {
92                Ok((parent, _)) => current = parent,
93                Err(_) => break,
94            }
95        }
96        for &node in overflow_stack.iter().rev() {
97            unsafe { Spec::top_down(BstDataMutRef::new_unchecked(node)) };
98        }
99        while inline_len > 0 {
100            inline_len -= 1;
101            unsafe {
102                Spec::top_down(BstDataMutRef::new_unchecked(
103                    *inline_stack[inline_len].assume_init_ref(),
104                ));
105            }
106        }
107
108        while let Ok((parent, node_direction)) = unsafe { internal_parent::<Spec, Data>(node) } {
109            if let Ok((_, parent_direction)) = unsafe { internal_parent::<Spec, Data>(parent) } {
110                if node_direction == parent_direction {
111                    unsafe { rotate::<Spec, Data>(parent) };
112                } else {
113                    unsafe { rotate::<Spec, Data>(node) };
114                }
115            }
116            unsafe { rotate::<Spec, Data>(node) };
117        }
118        unsafe { Spec::bottom_up(BstDataMutRef::new_unchecked(node)) };
119        current
120    }
121
122    /// Moves `node` to the root by propagating only the nodes involved in each rotation and
123    /// returns the previous root.
124    ///
125    /// # Safety
126    ///
127    /// `node` and every pointer reachable through its auxiliary-parent chain must refer to live
128    /// nodes of the same tree. Propagating an ancestor after its descendant must be valid for
129    /// `Spec`.
130    #[inline(always)]
131    pub unsafe fn splay_with_local_top_down<Spec, Data>(mut node: NodePtr<Spec>) -> NodePtr<Spec>
132    where
133        Spec: BstSpec<Data = Data, Parent = WithParent<Data>>,
134    {
135        let mut current = node;
136        unsafe { Spec::top_down(BstDataMutRef::new_unchecked(node)) };
137        while let Ok((parent, _)) = unsafe { internal_parent::<Spec, Data>(node) } {
138            match unsafe { internal_parent::<Spec, Data>(parent) } {
139                Ok((grandparent, _)) => {
140                    current = grandparent;
141                    unsafe {
142                        Spec::top_down(BstDataMutRef::new_unchecked(grandparent));
143                        Spec::top_down(BstDataMutRef::new_unchecked(parent));
144                        Spec::top_down(BstDataMutRef::new_unchecked(node));
145                        let node_direction = usize::from(parent.as_ref().child[1] == Some(node));
146                        let parent_direction =
147                            usize::from(grandparent.as_ref().child[1] == Some(parent));
148                        node.as_mut().parent.parent = grandparent.as_ref().parent.parent;
149                        if let Ok((mut ancestor, direction)) =
150                            internal_parent::<Spec, Data>(grandparent)
151                        {
152                            ancestor.as_mut().child[direction] = Some(node);
153                        }
154                        if node_direction == parent_direction {
155                            rotate_at::<Spec, Data>(parent, grandparent, parent_direction);
156                            rotate_at::<Spec, Data>(node, parent, node_direction);
157                            Spec::bottom_up(BstDataMutRef::new_unchecked(grandparent));
158                            Spec::bottom_up(BstDataMutRef::new_unchecked(parent));
159                        } else {
160                            rotate_at::<Spec, Data>(node, parent, node_direction);
161                            rotate_at::<Spec, Data>(node, grandparent, parent_direction);
162                            Spec::bottom_up(BstDataMutRef::new_unchecked(parent));
163                            Spec::bottom_up(BstDataMutRef::new_unchecked(grandparent));
164                        }
165                    }
166                }
167                Err(ancestor) => {
168                    current = parent;
169                    unsafe {
170                        Spec::top_down(BstDataMutRef::new_unchecked(parent));
171                        Spec::top_down(BstDataMutRef::new_unchecked(node));
172                        let direction = usize::from(parent.as_ref().child[1] == Some(node));
173                        node.as_mut().parent.parent = ancestor;
174                        rotate_at::<Spec, Data>(node, parent, direction);
175                        Spec::bottom_up(BstDataMutRef::new_unchecked(parent));
176                    }
177                }
178            }
179        }
180        unsafe { Spec::bottom_up(BstDataMutRef::new_unchecked(node)) };
181        current
182    }
183}
184
185pub fn rooted_heavy_order(
186    vertices_size: usize,
187    edges: &[(usize, usize)],
188) -> Vec<(usize, usize, bool)> {
189    if vertices_size == 0 {
190        return Vec::new();
191    }
192    let mut head = vec![usize::MAX; vertices_size];
193    let mut to = Vec::with_capacity(edges.len() * 2);
194    let mut next = Vec::with_capacity(edges.len() * 2);
195    for &(u, v) in edges {
196        to.push(v);
197        next.push(head[u]);
198        head[u] = to.len() - 1;
199        to.push(u);
200        next.push(head[v]);
201        head[v] = to.len() - 1;
202    }
203    let mut parent = vec![usize::MAX; vertices_size];
204    let mut stack = vec![0];
205    let mut order = Vec::with_capacity(vertices_size - 1);
206    parent[0] = 0;
207    while let Some(u) = stack.pop() {
208        let mut edge = head[u];
209        while edge != usize::MAX {
210            let v = to[edge];
211            if parent[v] == usize::MAX {
212                parent[v] = u;
213                order.push((v, u));
214                stack.push(v);
215            }
216            edge = next[edge];
217        }
218    }
219    let mut size = vec![1usize; vertices_size];
220    let mut heavy = vec![usize::MAX; vertices_size];
221    for &(child, parent) in order.iter().rev() {
222        size[parent] += size[child];
223        if heavy[parent] == usize::MAX || size[heavy[parent]] < size[child] {
224            heavy[parent] = child;
225        }
226    }
227    order
228        .into_iter()
229        .map(|(child, parent)| (child, parent, heavy[parent] == child))
230        .collect()
231}
232
233#[inline]
234pub fn splay<Spec, Data, Seeker>(
235    root: BstRoot<Spec>,
236    mut seeker: Seeker,
237) -> (Ordering, BstRoot<Spec>)
238where
239    Spec: BstSpec<Data = Data, Parent = WithNoParent<Data>>,
240    Seeker: BstSeeker<Spec = Spec>,
241{
242    let mut root = root;
243    let mut left_subtree = None;
244    let mut right_subtree = None;
245    let mut left_entry = &mut left_subtree;
246    let mut right_entry = &mut right_subtree;
247    let mut inline_stack = [None; 24];
248    let mut inline_len = 0;
249    let mut overflow_stack = vec![];
250
251    macro_rules! push_node {
252        ($node:expr) => {
253            if inline_len < inline_stack.len() {
254                inline_stack[inline_len] = Some($node);
255                inline_len += 1;
256            } else {
257                overflow_stack.push($node);
258            }
259        };
260    }
261
262    macro_rules! add {
263        (@left $node:ident) => {
264            *left_entry = Some($node.node);
265            push_node!($node.node);
266            left_entry = unsafe { &mut $node.node.as_mut().child[1] };
267        };
268        (@right $node:ident) => {
269            *right_entry = Some($node.node);
270            push_node!($node.node);
271            right_entry = unsafe { &mut $node.node.as_mut().child[0] };
272        };
273    }
274
275    let root_ordering = loop {
276        Spec::top_down(root.borrow_datamut());
277        match seeker.bst_seek(root.reborrow()) {
278            Ordering::Greater => {
279                let Some(mut child) = (unsafe { root.borrow_mut().left_mut().take() }) else {
280                    break Ordering::Greater;
281                };
282                Spec::top_down(child.borrow_datamut());
283                match seeker.bst_seek(child.reborrow()) {
284                    Ordering::Greater => {
285                        let Some(mut grandchild) =
286                            (unsafe { child.borrow_mut().left_mut().take() })
287                        else {
288                            add!(@right root);
289                            root = child;
290                            break Ordering::Greater;
291                        };
292                        Spec::top_down(grandchild.borrow_datamut());
293                        let child_right = unsafe { child.borrow_mut().right_mut().take() };
294                        if let Some(child_right) = child_right {
295                            unsafe { root.borrow_mut().left_mut().set(child_right) };
296                        }
297                        Spec::bottom_up(root.borrow_datamut());
298                        unsafe { child.borrow_mut().right_mut().set(root) };
299                        add!(@right child);
300                        root = grandchild;
301                    }
302                    Ordering::Equal => {
303                        add!(@right root);
304                        root = child;
305                        break Ordering::Equal;
306                    }
307                    Ordering::Less => {
308                        let Some(mut grandchild) =
309                            (unsafe { child.borrow_mut().right_mut().take() })
310                        else {
311                            add!(@right root);
312                            root = child;
313                            break Ordering::Less;
314                        };
315                        Spec::top_down(grandchild.borrow_datamut());
316                        add!(@right root);
317                        add!(@left child);
318                        root = grandchild;
319                    }
320                }
321            }
322            Ordering::Equal => break Ordering::Equal,
323            Ordering::Less => {
324                let Some(mut child) = (unsafe { root.borrow_mut().right_mut().take() }) else {
325                    break Ordering::Less;
326                };
327                Spec::top_down(child.borrow_datamut());
328                match seeker.bst_seek(child.reborrow()) {
329                    Ordering::Greater => {
330                        let Some(mut grandchild) =
331                            (unsafe { child.borrow_mut().left_mut().take() })
332                        else {
333                            add!(@left root);
334                            root = child;
335                            break Ordering::Greater;
336                        };
337                        Spec::top_down(grandchild.borrow_datamut());
338                        add!(@left root);
339                        add!(@right child);
340                        root = grandchild;
341                    }
342                    Ordering::Equal => {
343                        add!(@left root);
344                        root = child;
345                        break Ordering::Equal;
346                    }
347                    Ordering::Less => {
348                        let Some(mut grandchild) =
349                            (unsafe { child.borrow_mut().right_mut().take() })
350                        else {
351                            add!(@left root);
352                            root = child;
353                            break Ordering::Less;
354                        };
355                        Spec::top_down(grandchild.borrow_datamut());
356                        let child_left = unsafe { child.borrow_mut().left_mut().take() };
357                        if let Some(child_left) = child_left {
358                            unsafe { root.borrow_mut().right_mut().set(child_left) };
359                        }
360                        Spec::bottom_up(root.borrow_datamut());
361                        unsafe { child.borrow_mut().left_mut().set(root) };
362                        add!(@left child);
363                        root = grandchild;
364                    }
365                }
366            }
367        }
368    };
369
370    *left_entry = unsafe { root.borrow_mut().left_mut().take() }.map(|node| node.node);
371    *right_entry = unsafe { root.borrow_mut().right_mut().take() }.map(|node| node.node);
372    unsafe {
373        root.node.as_mut().child[0] = left_subtree;
374        root.node.as_mut().child[1] = right_subtree;
375        while let Some(node) = overflow_stack.pop() {
376            Spec::bottom_up(BstRoot::new(node).borrow_datamut());
377        }
378        while inline_len > 0 {
379            inline_len -= 1;
380            let node = inline_stack[inline_len].unwrap_unchecked();
381            Spec::bottom_up(BstRoot::new(node).borrow_datamut());
382        }
383    }
384    Spec::bottom_up(root.borrow_datamut());
385    (root_ordering, root)
386}
387
388#[inline]
389pub fn merge<Spec, Data>(
390    left: Option<BstRoot<Spec>>,
391    right: Option<BstRoot<Spec>>,
392) -> Option<BstRoot<Spec>>
393where
394    Spec: BstSpec<Data = Data, Parent = WithNoParent<Data>>,
395{
396    match (left, right) {
397        (None, None) => None,
398        (None, Some(root)) | (Some(root), None) => Some(root),
399        (Some(left), Some(mut right)) if right.reborrow().left().descend().is_err() => {
400            Spec::top_down(right.borrow_datamut());
401            unsafe { right.borrow_mut().left_mut().set(left) };
402            Spec::bottom_up(right.borrow_datamut());
403            Some(right)
404        }
405        (Some(left), Some(right)) => {
406            let (_, mut root) = splay(left, SeekRight::default());
407            unsafe { root.borrow_mut().right_mut().set(right) };
408            Spec::bottom_up(root.borrow_datamut());
409            Some(root)
410        }
411    }
412}
413
414#[inline]
415pub fn split<Spec, Data, Seeker>(
416    root: Option<BstRoot<Spec>>,
417    seeker: Seeker,
418    equal_side: EqualSide,
419) -> (Option<BstRoot<Spec>>, Option<BstRoot<Spec>>)
420where
421    Spec: BstSpec<Data = Data, Parent = WithNoParent<Data>>,
422    Seeker: BstSeeker<Spec = Spec>,
423{
424    let Some(root) = root else {
425        return (None, None);
426    };
427    let (ordering, mut root) = splay(root, seeker);
428    if equal_side.goes_left(ordering) {
429        let right = unsafe { root.borrow_mut().right_mut().take() };
430        Spec::bottom_up(root.borrow_datamut());
431        (Some(root), right)
432    } else {
433        let left = unsafe { root.borrow_mut().left_mut().take() };
434        Spec::bottom_up(root.borrow_datamut());
435        (left, Some(root))
436    }
437}
competitive/data_structure/splay_operations.rs

competitive/data_structure/
splay_operations.rs
