competitive/data_structure/splay_tree/
node.rs1use super::Allocator;
2use std::{
3 cmp::Ordering,
4 fmt::{self, Debug},
5 marker::PhantomData,
6 mem::{replace, swap},
7 ops::Bound,
8 ptr::NonNull,
9};
10
11pub trait SplaySpec: Sized {
12 type T;
13 fn has_bottom_up() -> bool {
14 false
15 }
16 fn top_down(_node: NodeRef<marker::DataMut<'_>, Self>) {}
17 fn bottom_up(_node: NodeRef<marker::DataMut<'_>, Self>) {}
18}
19
20pub trait SplaySeeker {
21 type S: SplaySpec;
22 fn splay_seek(&mut self, _node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering;
23}
24
25pub struct SeekLeft<S> {
26 _marker: PhantomData<fn() -> S>,
27}
28impl<S> Default for SeekLeft<S> {
29 fn default() -> Self {
30 Self {
31 _marker: PhantomData,
32 }
33 }
34}
35impl<S> SeekLeft<S> {
36 pub fn new() -> Self {
37 Self::default()
38 }
39}
40impl<S> SplaySeeker for SeekLeft<S>
41where
42 S: SplaySpec,
43{
44 type S = S;
45 fn splay_seek(&mut self, _node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
46 Ordering::Less
47 }
48}
49
50pub struct SeekRight<S> {
51 _marker: PhantomData<fn() -> S>,
52}
53impl<S> Default for SeekRight<S> {
54 fn default() -> Self {
55 Self {
56 _marker: PhantomData,
57 }
58 }
59}
60impl<S> SeekRight<S> {
61 pub fn new() -> Self {
62 Self::default()
63 }
64}
65impl<S> SplaySeeker for SeekRight<S>
66where
67 S: SplaySpec,
68{
69 type S = S;
70 fn splay_seek(&mut self, _node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
71 Ordering::Greater
72 }
73}
74
75pub struct Node<T> {
76 data: T,
77 left: Option<NonNull<Node<T>>>,
78 right: Option<NonNull<Node<T>>>,
79}
80
81impl<T> Node<T> {
82 pub fn new(data: T) -> Self {
83 Self {
84 data,
85 left: None,
86 right: None,
87 }
88 }
89}
90
91pub struct NodeRef<B, S>
92where
93 S: SplaySpec,
94{
95 node: NonNull<Node<S::T>>,
96 _marker: PhantomData<B>,
97}
98
99pub struct Root<S>
100where
101 S: SplaySpec,
102{
103 root: Option<NodeRef<marker::Owned, S>>,
104}
105
106impl<S> Default for Root<S>
107where
108 S: SplaySpec,
109{
110 fn default() -> Self {
111 Self { root: None }
112 }
113}
114
115impl<'a, S> Copy for NodeRef<marker::Immut<'a>, S>
116where
117 S: SplaySpec,
118 S::T: 'a,
119{
120}
121impl<'a, S> Clone for NodeRef<marker::Immut<'a>, S>
122where
123 S: SplaySpec,
124 S::T: 'a,
125{
126 fn clone(&self) -> Self {
127 *self
128 }
129}
130
131impl<B, S> NodeRef<B, S>
132where
133 S: SplaySpec,
134{
135 unsafe fn new_unchecked(node: NonNull<Node<S::T>>) -> Self {
136 Self {
137 node,
138 _marker: PhantomData,
139 }
140 }
141 fn reborrow(&self) -> NodeRef<marker::Immut<'_>, S> {
142 unsafe { NodeRef::new_unchecked(self.node) }
143 }
144 fn as_ptr(&self) -> *mut Node<S::T> {
145 self.node.as_ptr()
146 }
147}
148
149impl<S> NodeRef<marker::Owned, S>
150where
151 S: SplaySpec,
152{
153 pub fn new(node: NonNull<Node<S::T>>) -> Self {
154 unsafe { NodeRef::new_unchecked(node) }
155 }
156 pub unsafe fn from_data<A>(data: S::T, allocator: &mut A) -> Self
157 where
158 A: Allocator<Node<S::T>>,
159 {
160 Self::new(allocator.allocate(Node::new(data)))
161 }
162 pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, S> {
163 unsafe { NodeRef::new_unchecked(self.node) }
164 }
165 pub fn borrow_datamut(&mut self) -> NodeRef<marker::DataMut<'_>, S> {
166 unsafe { NodeRef::new_unchecked(self.node) }
167 }
168 pub fn into_dying(self) -> NodeRef<marker::Dying, S> {
169 unsafe { NodeRef::new_unchecked(self.node) }
170 }
171}
172
173impl<'a, S> NodeRef<marker::Immut<'a>, S>
174where
175 S: SplaySpec,
176 S::T: 'a,
177{
178 pub fn data(&self) -> &'a S::T {
179 unsafe { &(*self.as_ptr()).data }
180 }
181 pub fn left(&self) -> Option<Self> {
182 unsafe {
183 (*self.as_ptr())
184 .left
185 .map(|node| NodeRef::new_unchecked(node))
186 }
187 }
188 pub fn right(&self) -> Option<Self> {
189 unsafe {
190 (*self.as_ptr())
191 .right
192 .map(|node| NodeRef::new_unchecked(node))
193 }
194 }
195 pub fn traverse<F>(self, f: &mut F)
196 where
197 S::T: 'a,
198 F: FnMut(Self),
199 {
200 if let Some(left) = self.clone().left() {
201 left.traverse(f);
202 }
203 f(self);
204 if let Some(right) = self.clone().right() {
205 right.traverse(f);
206 }
207 }
208}
209
210impl<'a, S> NodeRef<marker::DataMut<'a>, S>
211where
212 S: SplaySpec,
213 S::T: 'a,
214{
215 pub fn data(&self) -> &'a S::T {
216 unsafe { &(*self.as_ptr()).data }
217 }
218 pub fn data_mut(&self) -> &'a mut S::T {
219 unsafe { &mut (*self.as_ptr()).data }
220 }
221 pub fn left(&self) -> Option<Self> {
222 unsafe {
223 (*self.as_ptr())
224 .left
225 .map(|node| NodeRef::new_unchecked(node))
226 }
227 }
228 pub fn right(&self) -> Option<Self> {
229 unsafe {
230 (*self.as_ptr())
231 .right
232 .map(|node| NodeRef::new_unchecked(node))
233 }
234 }
235 pub fn reverse(&self) {
236 unsafe {
237 let node = &mut (*self.as_ptr());
238 swap(&mut node.left, &mut node.right);
239 }
240 }
241}
242
243impl<'a, S> NodeRef<marker::Mut<'a>, S>
244where
245 S: SplaySpec,
246 S::T: 'a,
247{
248 pub fn data(self) -> &'a S::T {
249 unsafe { &(*self.as_ptr()).data }
250 }
251 pub fn data_mut(self) -> &'a mut S::T {
252 unsafe { &mut (*self.as_ptr()).data }
253 }
254 pub fn take_left(&mut self) -> Option<NodeRef<marker::Owned, S>> {
255 Some(NodeRef::new(unsafe { (*self.as_ptr()).left.take()? }))
256 }
257 pub fn take_right(&mut self) -> Option<NodeRef<marker::Owned, S>> {
258 Some(NodeRef::new(unsafe { (*self.as_ptr()).right.take()? }))
259 }
260 pub fn set_left(&mut self, node: Option<NodeRef<marker::Owned, S>>) {
261 unsafe { (*self.as_ptr()).left = node.map(|node| node.node) }
262 }
263 pub fn set_right(&mut self, node: Option<NodeRef<marker::Owned, S>>) {
264 unsafe { (*self.as_ptr()).right = node.map(|node| node.node) }
265 }
266}
267
268impl<S> NodeRef<marker::Dying, S>
269where
270 S: SplaySpec,
271{
272 pub unsafe fn into_inner(self) -> NonNull<Node<S::T>> {
273 let node = self.node;
274 unsafe {
275 debug_assert!((*node.as_ptr()).left.is_none());
276 debug_assert!((*node.as_ptr()).right.is_none());
277 }
278 node
279 }
280 pub unsafe fn into_data<A>(self, allocator: &mut A) -> S::T
281 where
282 A: Allocator<Node<S::T>>,
283 {
284 let Node { data, left, right } = allocator.deallocate(self.node);
285 debug_assert!(left.is_none());
286 debug_assert!(right.is_none());
287 data
288 }
289}
290
291impl<S> NodeRef<marker::Owned, S>
292where
293 S: SplaySpec,
294{
295 pub fn splay_by<Seeker>(self, mut seeker: Seeker) -> (Ordering, Self)
297 where
298 Seeker: SplaySeeker<S = S>,
299 {
300 let mut x = self;
301 let mut left_subtree: Option<NonNull<Node<S::T>>> = None;
303 let mut right_subtree: Option<NonNull<Node<S::T>>> = None;
304 let mut left_entry = &mut left_subtree;
305 let mut right_entry = &mut right_subtree;
306 let mut stack = vec![];
307
308 macro_rules! add {
309 (@left Some($ptr:ident)) => { add!(@inner Some($ptr.node), left_entry $ptr right); };
310 (@right Some($ptr:ident)) => { add!(@inner Some($ptr.node), right_entry $ptr left); };
311 (@inner $node:expr, $entry:ident $ptr:ident $dir:ident) => {
312 *$entry = $node;
313 if S::has_bottom_up() {
314 stack.push($ptr.node);
315 }
316 $entry = unsafe { &mut (*$entry.as_mut().unwrap().as_ptr()).$dir };
317 };
318 }
319
320 let root_ord = loop {
321 S::top_down(x.borrow_datamut());
322 match seeker.splay_seek(x.reborrow()) {
323 Ordering::Less => {
324 if let Some(mut y) = x.borrow_mut().take_left() {
325 S::top_down(y.borrow_datamut());
326 match seeker.splay_seek(y.reborrow()) {
327 Ordering::Less => {
328 if let Some(mut z) = y.borrow_mut().take_left() {
329 S::top_down(z.borrow_datamut());
330 x.borrow_mut().set_left(y.borrow_mut().take_right());
331 S::bottom_up(x.borrow_datamut());
332 y.borrow_mut().set_right(Some(x));
333 add!(@right Some(y));
334 x = z;
335 } else {
336 add!(@right Some(x));
337 x = y;
338 break Ordering::Less;
339 }
340 }
341 Ordering::Equal => {
342 add!(@right Some(x));
343 x = y;
344 break Ordering::Equal;
345 }
346 Ordering::Greater => {
347 if let Some(mut z) = y.borrow_mut().take_right() {
348 S::top_down(z.borrow_datamut());
349 add!(@right Some(x));
350 add!(@left Some(y));
351 x = z;
352 } else {
353 add!(@right Some(x));
354 x = y;
355 break Ordering::Greater;
356 }
357 }
358 }
359 } else {
360 break Ordering::Less;
361 }
362 }
363 Ordering::Equal => break Ordering::Equal,
364 Ordering::Greater => {
365 if let Some(mut y) = x.borrow_mut().take_right() {
366 S::top_down(y.borrow_datamut());
367 match seeker.splay_seek(y.reborrow()) {
368 Ordering::Less => {
369 if let Some(mut z) = y.borrow_mut().take_left() {
370 S::top_down(z.borrow_datamut());
371 add!(@left Some(x));
372 add!(@right Some(y));
373 x = z;
374 } else {
375 add!(@left Some(x));
376 x = y;
377 break Ordering::Less;
378 }
379 }
380 Ordering::Equal => {
381 add!(@left Some(x));
382 x = y;
383 break Ordering::Equal;
384 }
385 Ordering::Greater => {
386 if let Some(mut z) = y.borrow_mut().take_right() {
387 S::top_down(z.borrow_datamut());
388 x.borrow_mut().set_right(y.borrow_mut().take_left());
389 S::bottom_up(x.borrow_datamut());
390 y.borrow_mut().set_left(Some(x));
391 add!(@left Some(y));
392 x = z;
393 } else {
394 add!(@left Some(x));
395 x = y;
396 break Ordering::Greater;
397 }
398 }
399 }
400 } else {
401 break Ordering::Greater;
402 }
403 }
404 }
405 };
406 *left_entry = x.borrow_mut().take_left().map(|node| node.node);
407 *right_entry = x.borrow_mut().take_right().map(|node| node.node);
408 unsafe {
409 x.borrow_mut()
410 .set_left(left_subtree.map(|node| NodeRef::new_unchecked(node)));
411 x.borrow_mut()
412 .set_right(right_subtree.map(|node| NodeRef::new_unchecked(node)));
413 if S::has_bottom_up() {
414 while let Some(node) = stack.pop() {
415 S::bottom_up(NodeRef::new_unchecked(node));
416 }
417 }
418 }
419 S::bottom_up(x.borrow_datamut());
420 (root_ord, x)
421 }
422 pub fn insert_left(mut self, mut node: Self) -> Self {
423 if let Some(left) = self.borrow_mut().take_left() {
424 node.borrow_mut().set_left(Some(left));
425 S::bottom_up(self.borrow_datamut());
426 };
427 node.borrow_mut().set_right(Some(self));
428 S::bottom_up(node.borrow_datamut());
429 node
430 }
431 pub fn insert_right(mut self, mut node: Self) -> Self {
432 if let Some(right) = self.borrow_mut().take_right() {
433 node.borrow_mut().set_right(Some(right));
434 S::bottom_up(self.borrow_datamut());
435 }
436 node.borrow_mut().set_left(Some(self));
437 S::bottom_up(node.borrow_datamut());
438 node
439 }
440 pub fn insert_first(self, mut node: Self) -> Self {
441 node.borrow_mut().set_right(Some(self));
442 S::bottom_up(node.borrow_datamut());
443 node
444 }
445 pub fn insert_last(self, mut node: Self) -> Self {
446 node.borrow_mut().set_left(Some(self));
447 S::bottom_up(node.borrow_datamut());
448 node
449 }
450 pub fn merge(mut self, mut other: Self) -> Self {
451 if other.reborrow().left().is_none() {
452 S::top_down(other.borrow_datamut());
453 other.borrow_mut().set_left(Some(self));
454 S::bottom_up(other.borrow_datamut());
455 other
456 } else {
457 self = self.splay_by(SeekRight::new()).1;
458 self.borrow_mut().set_right(Some(other));
459 S::bottom_up(self.borrow_datamut());
460 self
461 }
462 }
463}
464
465impl<S> Root<S>
466where
467 S: SplaySpec,
468{
469 pub fn new(root: Option<NodeRef<marker::Owned, S>>) -> Self {
470 Self { root }
471 }
472 pub unsafe fn from_single_nodes(nodes: Vec<NodeRef<marker::Owned, S>>) -> Self {
473 unsafe { Self::from_single_nodes_inner(&nodes) }
474 }
475 unsafe fn from_single_nodes_inner(nodes: &[NodeRef<marker::Owned, S>]) -> Self {
476 if nodes.is_empty() {
477 Self::new(None)
478 } else {
479 let m = nodes.len() / 2;
480 let left = unsafe { Self::from_single_nodes_inner(&nodes[..m]) };
481 let right = unsafe { Self::from_single_nodes_inner(&nodes[m + 1..]) };
482 let mut node = NodeRef::new(nodes[m].node);
483 node.borrow_mut().set_left(left.root);
484 node.borrow_mut().set_right(right.root);
485 S::bottom_up(node.borrow_datamut());
486 Self::new(Some(node))
487 }
488 }
489 pub fn is_empty(&self) -> bool {
490 self.root.is_none()
491 }
492 pub fn root(&self) -> Option<NodeRef<marker::Immut<'_>, S>> {
493 Some(self.root.as_ref()?.reborrow())
494 }
495 pub fn root_data_mut(&mut self) -> Option<NodeRef<marker::DataMut<'_>, S>> {
496 Some(self.root.as_mut()?.borrow_datamut())
497 }
498 pub fn splay_by<Seeker>(&mut self, seeker: Seeker) -> Option<Ordering>
499 where
500 Seeker: SplaySeeker<S = S>,
501 {
502 let (ord, root) = self.root.take()?.splay_by(seeker);
503 self.root = Some(root);
504 Some(ord)
505 }
506 pub fn insert_left(&mut self, mut node: NodeRef<marker::Owned, S>) {
507 self.root = Some(match self.root.take() {
508 Some(root) => root.insert_left(node),
509 None => {
510 S::bottom_up(node.borrow_datamut());
511 node
512 }
513 });
514 }
515 pub fn insert_right(&mut self, mut node: NodeRef<marker::Owned, S>) {
516 self.root = Some(match self.root.take() {
517 Some(root) => root.insert_right(node),
518 None => {
519 S::bottom_up(node.borrow_datamut());
520 node
521 }
522 });
523 }
524 pub fn insert_first(&mut self, mut node: NodeRef<marker::Owned, S>) {
525 self.root = Some(match self.root.take() {
526 Some(root) => root.insert_first(node),
527 None => {
528 S::bottom_up(node.borrow_datamut());
529 node
530 }
531 });
532 }
533 pub fn insert_last(&mut self, mut node: NodeRef<marker::Owned, S>) {
534 self.root = Some(match self.root.take() {
535 Some(root) => root.insert_last(node),
536 None => {
537 S::bottom_up(node.borrow_datamut());
538 node
539 }
540 });
541 }
542 pub fn take_root(&mut self) -> Option<NodeRef<marker::Owned, S>> {
543 let mut root = self.root.take()?;
544 let right = root.borrow_mut().take_right();
545 self.root = root.borrow_mut().take_left();
546 self.append(&mut Self::new(right));
547 Some(root)
548 }
549 pub fn take_first(&mut self) -> Option<NodeRef<marker::Owned, S>> {
550 let mut root = self.root.take()?.splay_by(SeekLeft::new()).1;
551 let right = root.borrow_mut().take_right();
552 self.root = right;
553 Some(root)
554 }
555 pub fn take_last(&mut self) -> Option<NodeRef<marker::Owned, S>> {
556 let mut root = self.root.take()?.splay_by(SeekRight::new()).1;
557 let left = root.borrow_mut().take_left();
558 self.root = left;
559 Some(root)
560 }
561 pub fn split_left(&mut self) -> Option<NodeRef<marker::Owned, S>> {
562 let root = self.root.as_mut()?;
563 let left = root.borrow_mut().take_left();
564 S::bottom_up(root.borrow_datamut());
565 left
566 }
567 pub fn split_right(&mut self) -> Option<NodeRef<marker::Owned, S>> {
568 let root = self.root.as_mut()?;
569 let right = root.borrow_mut().take_right();
570 S::bottom_up(root.borrow_datamut());
571 right
572 }
573 pub fn split_left_eq(&mut self) -> Option<NodeRef<marker::Owned, S>> {
574 let right = self.split_right();
575 replace(&mut self.root, right)
576 }
577 pub fn split_right_eq(&mut self) -> Option<NodeRef<marker::Owned, S>> {
578 let left = self.split_left();
579 replace(&mut self.root, left)
580 }
581 pub fn append(&mut self, other: &mut Self) {
582 self.root = match (self.root.take(), other.root.take()) {
583 (Some(node), None) | (None, Some(node)) => Some(node),
584 (Some(left), Some(right)) => Some(left.merge(right)),
585 (None, None) => None,
586 }
587 }
588 pub fn range<Seeker>(&mut self, start: Bound<Seeker>, end: Bound<Seeker>) -> NodeRange<'_, S>
589 where
590 Seeker: SplaySeeker<S = S>,
591 {
592 if self.is_empty() {
593 return NodeRange::new(self);
594 }
595 let right = match end {
596 Bound::Included(seeker) => match self.splay_by(seeker).unwrap() {
597 Ordering::Greater | Ordering::Equal => self.split_right(),
598 Ordering::Less => self.split_right_eq(),
599 },
600 Bound::Excluded(seeker) => match self.splay_by(seeker).unwrap() {
601 Ordering::Greater => self.split_right(),
602 Ordering::Less | Ordering::Equal => self.split_right_eq(),
603 },
604 Bound::Unbounded => None,
605 };
606 if self.is_empty() {
607 return NodeRange::three_way(None, self, right);
608 }
609 let left = match start {
610 Bound::Included(seeker) => match self.splay_by(seeker).unwrap() {
611 Ordering::Less | Ordering::Equal => self.split_left(),
612 Ordering::Greater => self.split_left_eq(),
613 },
614 Bound::Excluded(seeker) => match self.splay_by(seeker).unwrap() {
615 Ordering::Less => self.split_left(),
616 Ordering::Greater | Ordering::Equal => self.split_left_eq(),
617 },
618 Bound::Unbounded => None,
619 };
620 NodeRange::three_way(left, self, right)
621 }
622}
623
624impl<S> Debug for Root<S>
625where
626 S: SplaySpec,
627 S::T: Debug,
628{
629 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
630 fn fmt_recurse<'a, S>(
631 node: NodeRef<marker::Immut<'a>, S>,
632 f: &mut fmt::Formatter<'_>,
633 ) -> fmt::Result
634 where
635 S: SplaySpec,
636 S::T: 'a + Debug,
637 {
638 write!(f, "[")?;
639 if let Some(left) = node.left() {
640 fmt_recurse(left, f)?;
641 }
642 node.data().fmt(f)?;
643 if let Some(right) = node.right() {
644 fmt_recurse(right, f)?;
645 }
646 write!(f, "]")?;
647 Ok(())
648 }
649 if let Some(root) = self.root.as_ref() {
650 let root = root.reborrow();
651 fmt_recurse(root, f)?;
652 }
653 Ok(())
654 }
655}
656
657pub struct NodeRange<'a, S>
658where
659 S: SplaySpec,
660 S::T: 'a,
661{
662 front: Root<S>,
663 back: Root<S>,
664 root: &'a mut Root<S>,
665}
666
667impl<'a, S> Debug for NodeRange<'a, S>
668where
669 S: SplaySpec,
670 S::T: 'a + Debug,
671{
672 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
673 f.debug_struct("NodeRange")
674 .field("front", &self.front)
675 .field("back", &self.back)
676 .field("root", &self.root)
677 .finish()
678 }
679}
680impl<'a, S> NodeRange<'a, S>
681where
682 S: SplaySpec,
683 S::T: 'a,
684{
685 pub fn new(root: &'a mut Root<S>) -> Self {
686 Self {
687 front: Default::default(),
688 back: Default::default(),
689 root,
690 }
691 }
692 pub fn three_way(
693 front: Option<NodeRef<marker::Owned, S>>,
694 root: &'a mut Root<S>,
695 back: Option<NodeRef<marker::Owned, S>>,
696 ) -> Self {
697 Self {
698 front: Root::new(front),
699 back: Root::new(back),
700 root,
701 }
702 }
703 pub fn next_checked(&mut self) -> Option<NodeRef<marker::DataMut<'a>, S>> {
704 let first = self.root.take_first()?;
705 let noderef = unsafe { NodeRef::new_unchecked(first.node) };
706 self.front.insert_last(first);
707 Some(noderef)
708 }
709 pub fn next_back_checked(&mut self) -> Option<NodeRef<marker::DataMut<'a>, S>> {
710 let last = self.root.take_last()?;
711 let noderef = unsafe { NodeRef::new_unchecked(last.node) };
712 self.back.insert_first(last);
713 Some(noderef)
714 }
715 pub fn root(&self) -> &Root<S> {
716 self.root
717 }
718 pub fn root_mut(&mut self) -> &mut Root<S> {
719 self.root
720 }
721 pub fn front(&self) -> &Root<S> {
722 &self.front
723 }
724 pub fn drop_rotate_left(mut self) {
725 self.root.append(&mut self.back);
726 self.root.append(&mut self.front);
727 }
728}
729impl<'a, S> Drop for NodeRange<'a, S>
730where
731 S: SplaySpec,
732 S::T: 'a,
733{
734 fn drop(&mut self) {
735 swap(self.root, &mut self.front);
736 self.root.append(&mut self.front);
737 self.root.append(&mut self.back);
738 }
739}
740
741pub mod marker {
742 use std::marker::PhantomData;
743
744 pub enum Owned {}
745 pub enum Dying {}
746 pub struct Immut<'a>(PhantomData<&'a ()>);
747 pub struct Mut<'a>(PhantomData<&'a mut ()>);
748 pub struct DataMut<'a>(PhantomData<&'a mut ()>);
749}