competitive_library/structure/
skew_heap_lazy.rs

1//! Skew Heap Lazy
2
3use std::mem::swap;
4#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
5pub struct Heap<T: Clone> {
6    pub cost: i64,
7    pub value: T,
8    pub lazy: Option<i64>,
9    pub left: Option<Box<Heap<T>>>,
10    pub right: Option<Box<Heap<T>>>,
11}
12
13impl<T: Clone> Heap<T> {
14    pub fn new(cost: i64, value: T) -> Option<Box<Heap<T>>> {
15        Some(Box::new(Heap {
16            cost,
17            value,
18            lazy: None,
19            left: None,
20            right: None,
21        }))
22    }
23}
24#[derive(Default, Clone)]
25pub struct SkewHeap<T: Clone> {
26    node: Option<Box<Heap<T>>>,
27}
28impl<T: Clone> SkewHeap<T> {
29    pub fn new() -> Self {
30        Self { node: None }
31    }
32
33    #[inline]
34    pub fn push(&mut self, cost: i64, value: T) {
35        SkewHeap::merge(&mut self.node, Heap::new(cost, value));
36    }
37    #[inline]
38    pub fn top(&self) -> Option<(i64, T)> {
39        Some((self.node.as_ref()?.cost, self.node.as_ref()?.value.clone()))
40    }
41    #[inline]
42    pub fn pop(&mut self) -> Option<(i64, T)> {
43        Self::propagate(&mut self.node);
44        let value = self.top()?;
45
46        let (mut left, right) = {
47            let mut tmp = self.node.take().unwrap();
48            (tmp.left.take(), tmp.right.take())
49        };
50        SkewHeap::merge(&mut left, right);
51        swap(&mut self.node, &mut left);
52
53        Some(value)
54    }
55
56    #[inline]
57    pub fn merge(a: &mut Option<Box<Heap<T>>>, mut b: Option<Box<Heap<T>>>) {
58        if a.is_none() {
59            swap(a, &mut b);
60            return;
61        }
62        if b.is_none() {
63            return;
64        }
65        Self::propagate(a);
66        Self::propagate(&mut b);
67
68        if a.as_ref().unwrap().cost > b.as_ref().unwrap().cost {
69            swap(a, &mut b);
70        }
71        SkewHeap::merge(&mut a.as_mut().unwrap().right, b);
72
73        let tmp = a.as_mut().unwrap();
74        swap(&mut tmp.left, &mut tmp.right);
75    }
76
77    #[inline]
78    pub fn add(&mut self, value: i64) {
79        self.node.as_mut().unwrap().lazy = Some(value);
80        Self::propagate(&mut self.node);
81    }
82    #[inline]
83    fn propagate(node: &mut Option<Box<Heap<T>>>) {
84        if let Some(n) = node.as_mut() {
85            if n.lazy.is_none() {
86                return;
87            }
88            if let Some(l) = n.left.as_mut() {
89                l.lazy = n.lazy;
90            }
91            if let Some(r) = n.right.as_mut() {
92                r.lazy = n.lazy;
93            }
94
95            n.cost += n.lazy.unwrap();
96            n.lazy = None;
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    #[test]
105    fn test_heap() {
106        let mut a = vec![SkewHeap::new(); 5];
107
108        for i in 0..30 {
109            a[i % 5].push(i as i64, 0);
110        }
111
112        for (i, e) in a.iter().enumerate() {
113            assert_eq!(e.top().unwrap().0, i as i64);
114        }
115
116        for i in 1..5 {
117            let buff = a[i].node.take();
118            SkewHeap::merge(&mut a[0].node, buff);
119        }
120
121        for i in 0..15 {
122            assert_eq!(a[0].pop().unwrap().0, i);
123        }
124        a[0].add(5);
125        for i in 15..30 {
126            assert_eq!(a[0].pop().unwrap().0, i + 5);
127        }
128    }
129}
competitive_library/structure/skew_heap_lazy.rs

competitive_library/structure/
skew_heap_lazy.rs
