haar_lib/ds/
lazy_skew_heap.rs

1//! 遅延加算付き融合可能ヒープ
2
3use crate::num::one_zero::Zero;
4use crate::trait_alias;
5use std::{mem::swap, ops::AddAssign};
6
7trait_alias!(
8    /// [`LazySkewHeap<T>`]が扱える型
9    Elem: Ord + Copy + Zero + AddAssign
10);
11
12#[derive(Debug, Clone)]
13struct Node<T> {
14    value: T,
15    lazy: T,
16    left: Option<Box<Node<T>>>,
17    right: Option<Box<Node<T>>>,
18}
19
20impl<T: Elem> Node<T> {
21    pub fn new(value: T) -> Self {
22        Self {
23            value,
24            lazy: T::zero(),
25            left: None,
26            right: None,
27        }
28    }
29
30    fn propagate(&mut self) {
31        self.value += self.lazy;
32        if let Some(left) = &mut self.left {
33            left.as_mut().lazy += self.lazy;
34        }
35        if let Some(right) = &mut self.right {
36            right.as_mut().lazy += self.lazy;
37        }
38        self.lazy = T::zero();
39    }
40
41    pub fn meld(&mut self, other: Option<Box<Node<T>>>) {
42        self.propagate();
43        if let Some(mut other) = other {
44            other.as_mut().propagate();
45
46            if self.value < other.value {
47                swap(self, other.as_mut());
48            }
49
50            match self.right.as_mut() {
51                Some(right) => right.meld(Some(other)),
52                None => self.right = Some(other),
53            }
54
55            swap(&mut self.left, &mut self.right);
56        }
57    }
58}
59
60/// 遅延加算付き融合可能ヒープ
61#[derive(Debug, Clone, Default)]
62pub struct LazySkewHeap<T> {
63    root: Option<Box<Node<T>>>,
64    size: usize,
65}
66
67impl<T: Elem> LazySkewHeap<T> {
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: LazySkewHeap<T>) {
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
