haar_lib/ds/
lazy_skew_heap.rs

1//! 遅延加算付き融合可能ヒープ
2
3use crate::num::one_zero::Zero;
4use std::{mem::swap, ops::AddAssign};
5
6#[derive(Debug, Clone)]
7struct Node<T> {
8    value: T,
9    lazy: T,
10    left: Option<Box<Node<T>>>,
11    right: Option<Box<Node<T>>>,
12}
13
14impl<T> Node<T>
15where
16    T: Ord + Copy + Zero + AddAssign,
17{
18    pub fn new(value: T) -> Self {
19        Self {
20            value,
21            lazy: T::zero(),
22            left: None,
23            right: None,
24        }
25    }
26
27    fn propagate(&mut self) {
28        self.value += self.lazy;
29        if let Some(left) = &mut self.left {
30            left.as_mut().lazy += self.lazy;
31        }
32        if let Some(right) = &mut self.right {
33            right.as_mut().lazy += self.lazy;
34        }
35        self.lazy = T::zero();
36    }
37
38    pub fn meld(&mut self, other: Option<Box<Self>>) {
39        self.propagate();
40        if let Some(mut other) = other {
41            other.as_mut().propagate();
42
43            if self.value < other.value {
44                swap(self, other.as_mut());
45            }
46
47            match self.right.as_mut() {
48                Some(right) => right.meld(Some(other)),
49                None => self.right = Some(other),
50            }
51
52            swap(&mut self.left, &mut self.right);
53        }
54    }
55}
56
57/// 遅延加算付き融合可能ヒープ
58#[derive(Debug, Clone, Default)]
59pub struct LazySkewHeap<T> {
60    root: Option<Box<Node<T>>>,
61    size: usize,
62}
63
64impl<T> LazySkewHeap<T>
65where
66    T: Ord + Copy + Zero + AddAssign,
67{
68    /// 空の[`LazySkewHeap`]を生成する。
69    pub fn new() -> Self {
70        Self {
71            root: None,
72            size: 0,
73        }
74    }
75
76    /// 他の[`LazySkewHeap<T>`]を融合する。
77    pub fn meld(&mut self, other: Self) {
78        self.size += other.size;
79        match self.root.as_mut() {
80            None => self.root = other.root,
81            Some(root) => root.meld(other.root),
82        }
83    }
84
85    /// 値`value`を挿入する。
86    pub fn push(&mut self, value: T) {
87        self.size += 1;
88        let t = Some(Box::new(Node::new(value)));
89        match self.root.as_mut() {
90            None => self.root = t,
91            Some(root) => root.meld(t),
92        }
93    }
94
95    /// ヒープの最大値を返す。
96    pub fn peek(&self) -> Option<&T> {
97        self.root.as_ref().map(|x| &x.value)
98    }
99
100    /// 最大値をヒープから削除して、その値を返す。
101    pub fn pop(&mut self) -> Option<T> {
102        match self.root.take() {
103            None => None,
104            Some(root) => {
105                self.size -= 1;
106
107                let Node {
108                    value: x,
109                    left,
110                    right,
111                    ..
112                } = *root;
113                match left {
114                    None => self.root = right,
115                    Some(mut left) => {
116                        left.meld(right);
117                        self.root = Some(left);
118                    }
119                }
120
121                Some(x)
122            }
123        }
124    }
125
126    /// ヒープの全要素に値`value`を加算する。
127    pub fn add(&mut self, value: T) {
128        if let Some(root) = self.root.as_mut() {
129            root.lazy += value;
130            root.propagate();
131        }
132    }
133
134    /// ヒープに含まれている値の個数を返す。
135    pub fn len(&self) -> usize {
136        self.size
137    }
138
139    /// ヒープが空ならば`true`を返す。
140    pub fn is_empty(&self) -> bool {
141        self.size == 0
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use rand::Rng;
149    use std::collections::BinaryHeap;
150
151    #[test]
152    fn test() {
153        let mut heap = LazySkewHeap::<u32>::new();
154        let mut bheap = BinaryHeap::<u32>::new();
155
156        let mut heap2 = LazySkewHeap::<u32>::new();
157        let mut bheap2 = BinaryHeap::<u32>::new();
158
159        let mut rng = rand::thread_rng();
160
161        for _ in 0..100 {
162            let x = rng.gen_range(0..10000);
163            heap.push(x);
164            bheap.push(x);
165        }
166
167        let x = rng.gen_range(0..10000);
168        heap.add(x);
169        bheap = bheap.into_iter().map(|a| a + x).collect::<BinaryHeap<_>>();
170
171        for _ in 0..100 {
172            let x = rng.gen::<u32>();
173            heap2.push(x);
174            bheap2.push(x);
175        }
176
177        heap.meld(heap2);
178
179        while let Some(x) = bheap2.pop() {
180            bheap.push(x);
181        }
182
183        while let (Some(x), Some(y)) = (heap.pop(), bheap.pop()) {
184            assert_eq!(x, y);
185        }
186    }
187}
haar_lib/ds/lazy_skew_heap.rs

haar_lib/ds/
lazy_skew_heap.rs
