nekolib/ds/
removable_heap.rs

1//! 削除可能ヒープ。
2
3use std::collections::BinaryHeap;
4
5/// 削除可能ヒープ。
6///
7/// # Idea
8/// ヒープを二つ使う。
9///
10/// # Implementation notes
11/// `peek()` を `&self` にするため、`remove()` と `pop()` の直後はとりあえず
12/// `force()` している。
13///
14/// # Examples
15/// ```
16/// use nekolib::ds::RemovableHeap;
17///
18/// let mut heap = RemovableHeap::new();
19/// heap.push(2);
20/// heap.push(1);
21/// heap.push(3);
22/// heap.push(0);
23/// assert_eq!(heap.peek(), Some(&3)); // {0, 1, 2, 3}
24///
25/// heap.remove(2);                    // {0, 1, 3}
26/// assert_eq!(heap.pop(), Some(3));   // {0, 1}
27/// assert_eq!(heap.peek(), Some(&1));
28/// heap.remove(1);                    // {0}
29/// assert_eq!(heap.peek(), Some(&0));
30/// assert_eq!(heap.pop(), Some(0));   // {}
31/// assert!(heap.is_empty());
32/// ```
33///
34/// 削除する要素がヒープに入っていないとき、こまる。
35/// ```
36/// use nekolib::ds::RemovableHeap;
37///
38/// let mut heap = RemovableHeap::new();
39/// heap.push(1);
40/// heap.remove(2);
41/// heap.push(2);
42/// heap.push(3);                      // {1, 2, 3}?
43/// assert_eq!(heap.pop(), Some(3));   // seems good?
44/// assert_ne!(heap.peek(), Some(&2)); // but...
45/// assert_eq!(heap.peek(), Some(&1));
46/// ```
47///
48/// 未来の push 操作を打ち消すというわけでもない。
49/// ```
50/// use nekolib::ds::RemovableHeap;
51///
52/// let mut heap = RemovableHeap::new();
53/// heap.push(1);
54/// heap.push(2);
55/// heap.remove(100);
56/// heap.remove(3);
57/// heap.push(101);
58/// heap.push(3);                      // {1, 2, 101}?
59/// assert_eq!(heap.pop(), Some(101)); // seems good, so far
60/// assert_ne!(heap.peek(), Some(&2)); // but...
61/// assert_eq!(heap.peek(), Some(&3));
62/// ```
63#[derive(Clone)]
64pub struct RemovableHeap<T> {
65    alive: BinaryHeap<T>,
66    dead: BinaryHeap<T>,
67    len: usize,
68}
69
70impl<T: Ord> RemovableHeap<T> {
71    /// 空のヒープで初期化する。
72    pub fn new() -> Self {
73        Self {
74            alive: BinaryHeap::new(),
75            dead: BinaryHeap::new(),
76            len: 0,
77        }
78    }
79
80    /// 要素を追加する。
81    pub fn push(&mut self, item: T) {
82        self.len += 1;
83        self.alive.push(item);
84    }
85
86    /// 要素を削除する。
87    ///
88    /// # Notes
89    /// 削除処理を遅延させて行うため、`&T` ではなく `T` であることに注意。
90    ///
91    /// また、`item` はその時点でヒープに入っている要素のいずれかと等しい必要がある。
92    /// 「未来に追加される要素を打ち消す」あるいは「単に無視する」という仕様は、
93    /// どちらも操作順しだいでどちらもうまくいかなくなるはず。
94    ///
95    /// `BTreeSet` を使えば判定できるが、それなら最初から `BTreeSet` でやればよい。
96    pub fn remove(&mut self, item: T) {
97        self.len -= 1;
98        self.dead.push(item);
99        self.force();
100    }
101
102    /// 最大値を取り出す。
103    pub fn pop(&mut self) -> Option<T> {
104        self.force();
105        let res = self.alive.pop();
106        if res.is_some() {
107            self.len -= 1;
108            self.force();
109        }
110        res
111    }
112
113    /// 最大値を取得する。
114    pub fn peek(&self) -> Option<&T> { self.alive.peek() }
115
116    /// 空のとき `true` を返す。
117    pub fn is_empty(&self) -> bool { self.len == 0 }
118
119    /// 要素数を返す。
120    pub fn len(&self) -> usize { self.len }
121
122    fn force(&mut self) {
123        while let (Some(x), Some(y)) = (self.alive.peek(), self.dead.peek()) {
124            if x != y {
125                break;
126            }
127            self.alive.pop();
128            self.dead.pop();
129        }
130    }
131}
132
133#[test]
134fn test() {
135    let mut heap = RemovableHeap::new();
136    assert_eq!(heap.pop(), None);
137    heap.push(1); // {1}
138    assert_eq!(heap.peek(), Some(&1));
139    assert_eq!(heap.len(), 1);
140    heap.push(2); // {1, 2}
141    assert_eq!(heap.peek(), Some(&2));
142    assert_eq!(heap.len(), 2);
143    heap.remove(1); // {2}
144    assert_eq!(heap.len(), 1);
145    assert_eq!(heap.pop(), Some(2)); // {}
146    assert_eq!(heap.len(), 0);
147    assert_eq!(heap.pop(), None); // {}
148    assert_eq!(heap.len(), 0);
149
150    heap.push(2); // {2}
151    heap.push(1); // {1, 2}
152    assert_eq!(heap.len(), 2);
153    heap.remove(2); // {1}
154    assert_eq!(heap.len(), 1);
155    assert_eq!(heap.pop(), Some(1)); // {}
156    assert_eq!(heap.len(), 0);
157}
nekolib/ds/removable_heap.rs

nekolib/ds/
removable_heap.rs
