1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Skew Heap Lazy

use std::mem::swap;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct Heap<T: Clone> {
    pub cost: i64,
    pub value: T,
    pub lazy: Option<i64>,
    pub left: Option<Box<Heap<T>>>,
    pub right: Option<Box<Heap<T>>>,
}

impl<T: Clone> Heap<T> {
    pub fn new(cost: i64, value: T) -> Option<Box<Heap<T>>> {
        Some(Box::new(Heap {
            cost,
            value,
            lazy: None,
            left: None,
            right: None,
        }))
    }
}
#[derive(Default, Clone)]
pub struct SkewHeap<T: Clone> {
    node: Option<Box<Heap<T>>>,
}
impl<T: Clone> SkewHeap<T> {
    pub fn new() -> Self {
        Self { node: None }
    }

    #[inline]
    pub fn push(&mut self, cost: i64, value: T) {
        SkewHeap::merge(&mut self.node, Heap::new(cost, value));
    }
    #[inline]
    pub fn top(&self) -> Option<(i64, T)> {
        Some((self.node.as_ref()?.cost, self.node.as_ref()?.value.clone()))
    }
    #[inline]
    pub fn pop(&mut self) -> Option<(i64, T)> {
        Self::propagate(&mut self.node);
        let value = self.top()?;

        let (mut left, right) = {
            let mut tmp = self.node.take().unwrap();
            (tmp.left.take(), tmp.right.take())
        };
        SkewHeap::merge(&mut left, right);
        swap(&mut self.node, &mut left);

        Some(value)
    }

    #[inline]
    pub fn merge(a: &mut Option<Box<Heap<T>>>, mut b: Option<Box<Heap<T>>>) {
        if a.is_none() {
            swap(a, &mut b);
            return;
        }
        if b.is_none() {
            return;
        }
        Self::propagate(a);
        Self::propagate(&mut b);

        if a.as_ref().unwrap().cost > b.as_ref().unwrap().cost {
            swap(a, &mut b);
        }
        SkewHeap::merge(&mut a.as_mut().unwrap().right, b);

        let tmp = a.as_mut().unwrap();
        swap(&mut tmp.left, &mut tmp.right);
    }

    #[inline]
    pub fn add(&mut self, value: i64) {
        self.node.as_mut().unwrap().lazy = Some(value);
        Self::propagate(&mut self.node);
    }
    #[inline]
    fn propagate(node: &mut Option<Box<Heap<T>>>) {
        if let Some(n) = node.as_mut() {
            if n.lazy.is_none() {
                return;
            }
            if let Some(l) = n.left.as_mut() {
                l.lazy = n.lazy;
            }
            if let Some(r) = n.right.as_mut() {
                r.lazy = n.lazy;
            }

            n.cost += n.lazy.unwrap();
            n.lazy = None;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_heap() {
        let mut a = vec![SkewHeap::new(); 5];

        for i in 0..30 {
            a[i % 5].push(i as i64, 0);
        }

        for (i, e) in a.iter().enumerate() {
            assert_eq!(e.top().unwrap().0, i as i64);
        }

        for i in 1..5 {
            let buff = a[i].node.take();
            SkewHeap::merge(&mut a[0].node, buff);
        }

        for i in 0..15 {
            assert_eq!(a[0].pop().unwrap().0, i);
        }
        a[0].add(5);
        for i in 15..30 {
            assert_eq!(a[0].pop().unwrap().0, i + 5);
        }
    }
}