nekolib/ds/
skew_heap.rs

1pub struct SkewHeap<T> {
2    root: Link<T>,
3    len: usize,
4}
5
6type Link<T> = Option<Box<Node<T>>>;
7
8struct Node<T> {
9    elem: T,
10    left: Link<T>,
11    right: Link<T>,
12}
13
14impl<T> Node<T> {
15    pub fn new(elem: T) -> Self { Self { elem, left: None, right: None } }
16}
17
18impl<T: Ord> SkewHeap<T> {
19    pub fn new() -> Self { SkewHeap { root: None, len: 0 } }
20
21    pub fn len(&self) -> usize { self.len }
22    pub fn is_empty(&self) -> bool { self.len == 0 }
23    pub fn clear(&mut self) { while self.pop().is_some() {} }
24
25    pub fn peek(&self) -> Option<&T> {
26        self.root.as_ref().map(|node| &node.elem)
27    }
28
29    pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T>> {
30        if self.is_empty() {
31            None
32        } else {
33            Some(PeekMut { heap: self, tainted: false })
34        }
35    }
36
37    pub fn push(&mut self, elem: T) {
38        self.meld(Self { root: Some(Box::new(Node::new(elem))), len: 1 });
39    }
40
41    pub fn pop(&mut self) -> Option<T> {
42        self.root.take().map(|node| {
43            self.len -= 1;
44            self.root = Self::meld_internal(node.left, node.right);
45            node.elem
46        })
47    }
48
49    pub fn meld(&mut self, other: Self) {
50        self.len += other.len;
51        self.root = Self::meld_internal(self.root.take(), other.root);
52    }
53
54    fn meld_internal(left: Link<T>, right: Link<T>) -> Link<T> {
55        match (left, right) {
56            (left, None) => left,
57            (None, right) => right,
58            (Some(mut left), Some(right)) if left.elem >= right.elem => {
59                std::mem::swap(&mut left.left, &mut left.right);
60                left.left = Self::meld_internal(left.left.take(), Some(right));
61                Some(left)
62            }
63            (left, right) => Self::meld_internal(right, left),
64        }
65    }
66
67    fn fix_root(&mut self) {
68        let root = &self.root.as_ref().unwrap().elem;
69        let left = &self.root.as_ref().unwrap().left;
70        let right = &self.root.as_ref().unwrap().right;
71        if left.as_ref().map(|left| &left.elem > root).unwrap_or(false)
72            || right.as_ref().map(|right| &right.elem > root).unwrap_or(false)
73        {
74            let tmp = self.pop().unwrap();
75            self.push(tmp);
76        }
77    }
78}
79
80impl<T> SkewHeap<T> {
81    #[cfg(test)]
82    fn push_with_dir(&mut self, elem: T, dir: &[bool]) {
83        let mut from = &mut self.root;
84        for &is_right in dir.iter().chain(std::iter::once(&false)) {
85            if let Some(v) = from {
86                from = if is_right { &mut v.right } else { &mut v.left };
87            } else if from.is_none() {
88                *from = Some(Box::new(Node::new(elem)));
89                return;
90            }
91        }
92    }
93
94    #[cfg(test)]
95    fn in_order(&self) -> impl Iterator<Item = &T> {
96        fn dfs<'a, T>(v: &'a Link<T>, vec: &mut Vec<&'a T>) {
97            let v = match v {
98                Some(v) => v,
99                None => return,
100            };
101            dfs(&v.left, vec);
102            vec.push(&v.elem);
103            dfs(&v.right, vec);
104        }
105
106        let mut res = vec![];
107        dfs(&self.root, &mut res);
108        res.into_iter()
109    }
110
111    fn bfs_order(&self) -> impl Iterator<Item = (&T, Option<&T>, Option<&T>)> {
112        use std::collections::VecDeque;
113
114        let mut res = vec![];
115        let mut q = VecDeque::new();
116
117        q.extend(self.root.as_ref());
118        while let Some(v) = q.pop_front() {
119            let left = v.left.as_ref();
120            let right = v.right.as_ref();
121            res.push((&v.elem, left.map(|o| &o.elem), right.map(|o| &o.elem)));
122            q.extend(left.into_iter().chain(right));
123        }
124
125        res.into_iter()
126    }
127}
128
129impl<T: Ord> Extend<T> for SkewHeap<T> {
130    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
131        for item in iter {
132            self.push(item);
133        }
134    }
135}
136
137use std::iter::FromIterator;
138
139impl<T: Ord> FromIterator<T> for SkewHeap<T> {
140    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
141        let mut heap = Self::new();
142        heap.extend(iter);
143        heap
144    }
145}
146
147pub struct IntoIter<T>(SkewHeap<T>);
148
149impl<T: Ord> IntoIterator for SkewHeap<T> {
150    type IntoIter = IntoIter<T>;
151    type Item = T;
152
153    fn into_iter(self) -> Self::IntoIter { IntoIter(self) }
154}
155
156impl<T: Ord> Iterator for IntoIter<T> {
157    type Item = T;
158
159    fn next(&mut self) -> Option<Self::Item> { self.0.pop() }
160
161    fn size_hint(&self) -> (usize, Option<usize>) {
162        (self.0.len, Some(self.0.len))
163    }
164}
165
166impl<T: Ord> ExactSizeIterator for IntoIter<T> {
167    fn len(&self) -> usize { self.0.len }
168}
169
170pub struct PeekMut<'a, T: 'a + Ord> {
171    heap: &'a mut SkewHeap<T>,
172    tainted: bool,
173}
174
175use std::ops::{Deref, DerefMut};
176
177impl<T: Ord> Drop for PeekMut<'_, T> {
178    fn drop(&mut self) {
179        if self.tainted {
180            self.heap.fix_root();
181        }
182    }
183}
184
185impl<T: Ord> Deref for PeekMut<'_, T> {
186    type Target = T;
187
188    fn deref(&self) -> &T { &self.heap.root.as_ref().unwrap().elem }
189}
190
191impl<T: Ord> DerefMut for PeekMut<'_, T> {
192    fn deref_mut(&mut self) -> &mut T {
193        self.tainted = true;
194        &mut self.heap.root.as_mut().unwrap().elem
195    }
196}
197
198impl<'a, T: Ord> PeekMut<'a, T> {
199    pub fn pop(mut self) -> T {
200        let value = self.heap.pop().unwrap();
201        self.tainted = false;
202        value
203    }
204}
205
206impl<T: Ord> Default for SkewHeap<T> {
207    fn default() -> Self { Self::new() }
208}
209
210use std::fmt;
211
212impl<T: fmt::Debug> fmt::Debug for SkewHeap<T> {
213    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
214        f.debug_list().entries(self.bfs_order().map(|(x, ..)| x)).finish()
215    }
216}
217
218#[cfg(test)]
219mod test {
220    use super::SkewHeap;
221
222    #[test]
223    fn push_with_dir() {
224        let mut tr = SkewHeap::new();
225        tr.push_with_dir(1, &[]);
226        tr.push_with_dir(4, &[false]);
227        tr.push_with_dir(23, &[true]);
228        tr.push_with_dir(63, &[false, false]);
229        tr.push_with_dir(24, &[false, true]);
230        tr.push_with_dir(44, &[true, false]);
231        tr.push_with_dir(28, &[true, true]);
232
233        assert_eq!(tr.in_order().copied().collect::<Vec<_>>(), [
234            63, 4, 24, 1, 44, 23, 28
235        ]);
236    }
237
238    #[test]
239    fn meld() {
240        let mut left = SkewHeap::new();
241        left.push_with_dir(-1, &[]);
242        left.push_with_dir(-4, &[false]);
243        left.push_with_dir(-23, &[true]);
244        left.push_with_dir(-63, &[false, false]);
245        left.push_with_dir(-24, &[false, true]);
246        left.push_with_dir(-44, &[true, false]);
247        left.push_with_dir(-28, &[true, true]);
248
249        assert_eq!(left.in_order().copied().collect::<Vec<_>>(), [
250            -63, -4, -24, -1, -44, -23, -28
251        ]);
252
253        let mut right = SkewHeap::new();
254        right.push_with_dir(-13, &[]);
255        right.push_with_dir(-17, &[false]);
256        right.push_with_dir(-99, &[true]);
257        right.push_with_dir(-57, &[false, false]);
258        right.push_with_dir(-49, &[false, true]);
259        right.push_with_dir(-105, &[true, false]);
260        right.push_with_dir(-201, &[true, true]);
261
262        assert_eq!(right.in_order().copied().collect::<Vec<_>>(), [
263            -57, -17, -49, -13, -105, -99, -201
264        ]);
265
266        left.meld(right);
267        let melded: Vec<_> = left
268            .bfs_order()
269            .map(|(&v, l, r)| (v, l.copied(), r.copied()))
270            .collect();
271        assert_eq!(melded, [
272            (-1, Some(-13), Some(-4)),
273            (-13, Some(-23), Some(-17)),
274            (-4, Some(-63), Some(-24)),
275            (-23, Some(-28), Some(-44)),
276            (-17, Some(-57), Some(-49)),
277            (-63, None, None),
278            (-24, None, None),
279            (-28, Some(-99), None),
280            (-44, None, None),
281            (-57, None, None),
282            (-49, None, None),
283            (-99, Some(-105), Some(-201)),
284            (-105, None, None),
285            (-201, None, None),
286        ]);
287    }
288
289    #[test]
290    fn heap_sort() {
291        let mut a = vec![2, 9, 7, 3, 1, 6, 3, 5, 4, 8];
292        let mut heap = SkewHeap::new();
293        heap.extend(a.iter().copied());
294        a.sort_unstable();
295        a.reverse();
296        let mut b = vec![];
297        while let Some(x) = heap.pop() {
298            b.push(x);
299        }
300        assert_eq!(a, b);
301
302        assert_eq!(heap.len(), 0);
303        assert_eq!(heap.pop(), None);
304        assert_eq!(heap.pop(), None);
305        assert_eq!(heap.len(), 0);
306    }
307
308    #[test]
309    fn into_iter() {
310        let heap: SkewHeap<_> =
311            [2, 5, 1, 3, 4, 6, 6, 1].iter().copied().collect();
312        assert_eq!(heap.len(), 8);
313
314        let mut iter = heap.into_iter();
315        assert_eq!(iter.size_hint(), (8, Some(8)));
316        assert_eq!(iter.len(), 8);
317        assert_eq!(iter.next(), Some(6));
318        assert_eq!(iter.size_hint(), (7, Some(7)));
319        assert_eq!(iter.len(), 7);
320        assert_eq!(iter.next(), Some(6));
321        assert_eq!(iter.next(), Some(5));
322        assert_eq!(iter.len(), 5);
323        assert_eq!(iter.next(), Some(4));
324        assert_eq!(iter.len(), 4);
325        assert_eq!(iter.next(), Some(3));
326        assert_eq!(iter.len(), 3);
327        assert_eq!(iter.next(), Some(2));
328        assert_eq!(iter.len(), 2);
329        assert_eq!(iter.next(), Some(1));
330        assert_eq!(iter.len(), 1);
331        assert_eq!(iter.next(), Some(1));
332        assert_eq!(iter.len(), 0);
333        assert_eq!(iter.next(), None);
334        assert_eq!(iter.len(), 0);
335    }
336
337    #[test]
338    fn peek() {
339        use std::cmp::Reverse;
340
341        let mut heap: SkewHeap<_> =
342            [2, 5, 3, 1, 6, 2].iter().map(|&x| Reverse(x)).collect();
343
344        assert_eq!(heap.peek(), Some(&Reverse(1)));
345        heap.peek_mut().unwrap().0 -= 1; // [0, 2, 2, 3, 5, 6]
346        assert_eq!(heap.peek(), Some(&Reverse(0)));
347        heap.peek_mut().unwrap().0 += 4; // [2, 2, 3, 4, 5, 6]
348        assert_eq!(heap.peek(), Some(&Reverse(2)));
349        assert_eq!(heap.peek_mut().unwrap().pop().0, 2); // [2, 3, 4, 5, 6]
350        assert_eq!(heap.peek(), Some(&Reverse(2)));
351        assert_eq!(heap.peek_mut().unwrap().0, 2);
352        assert_eq!(heap.peek_mut().unwrap().pop().0, 2); // [3, 4, 5, 6]
353        assert_eq!(heap.peek(), Some(&Reverse(3)));
354        assert_eq!(heap.pop(), Some(Reverse(3))); // [4, 5, 6]
355        assert_eq!(heap.peek(), Some(&Reverse(4)));
356
357        heap.peek_mut().unwrap().0 += 10; // [5, 6, 14]
358        assert_eq!(heap.peek(), Some(&Reverse(5)));
359
360        heap.push(Reverse(8)); // [5, 6, 8, 14]
361        assert_eq!(heap.peek(), Some(&Reverse(5)));
362        assert_eq!(heap.pop(), Some(Reverse(5))); // [6, 8, 14]
363        assert_eq!(heap.pop(), Some(Reverse(6))); // [8, 14]
364        assert_eq!(heap.pop(), Some(Reverse(8))); // [14]
365        assert_eq!(heap.peek_mut().unwrap().pop().0, 14);
366        assert!(heap.peek_mut().is_none());
367        assert!(heap.peek().is_none());
368    }
369}
nekolib/ds/skew_heap.rs

nekolib/ds/
skew_heap.rs
