haar_lib/ds/
interval_heap.rs

1//! 最大値と最小値を得られるヒープ
2
3/// 最大値と最小値を得られるヒープ
4#[derive(Clone, Debug, Default)]
5pub struct IntervalHeap<T> {
6    data: Vec<T>,
7}
8
9#[inline]
10fn left(i: usize) -> usize {
11    (i + 1) * 2 - (i & 1)
12}
13
14#[inline]
15fn right(i: usize) -> usize {
16    left(i) + 2
17}
18
19#[inline]
20fn min(i: usize) -> usize {
21    i | 1
22}
23
24#[inline]
25fn max(i: usize) -> usize {
26    i & !1
27}
28
29#[inline]
30fn parent(i: usize) -> usize {
31    (((i - 2) >> 2) << 1) | (i & 1)
32}
33
34impl<T: Ord> IntervalHeap<T> {
35    #[inline]
36    fn min_index(&self) -> usize {
37        self.data.len().saturating_sub(1).min(1)
38    }
39
40    #[inline]
41    fn max_index(&self) -> usize {
42        0
43    }
44
45    fn top_down(&mut self, mut k: usize) -> usize {
46        let n = self.data.len();
47
48        if (k & 1) == 1 {
49            while left(k) < n {
50                let c = if n <= right(k) || self.data[left(k)] < self.data[right(k)] {
51                    left(k)
52                } else {
53                    right(k)
54                };
55
56                if c < n && self.data[c] < self.data[k] {
57                    self.data.swap(c, k);
58                    k = c;
59                } else {
60                    break;
61                }
62            }
63        } else {
64            while left(k) < n {
65                let c = if n <= right(k) || self.data[left(k)] > self.data[right(k)] {
66                    left(k)
67                } else {
68                    right(k)
69                };
70
71                if c < n && self.data[c] > self.data[k] {
72                    self.data.swap(c, k);
73                    k = c;
74                } else {
75                    break;
76                }
77            }
78        }
79
80        k
81    }
82
83    fn bottom_up(&mut self, mut k: usize) {
84        if min(k) < self.data.len() && self.data[max(k)] < self.data[min(k)] {
85            self.data.swap(max(k), min(k));
86            k ^= 1;
87        }
88
89        let root = 1;
90        while root < k {
91            let p = max(parent(k));
92            if self.data[k] <= self.data[p] {
93                break;
94            }
95            self.data.swap(p, k);
96            k = p;
97        }
98
99        while root < k {
100            let p = min(parent(k));
101            if self.data[k] >= self.data[p] {
102                break;
103            }
104            self.data.swap(p, k);
105            k = p;
106        }
107    }
108
109    /// 空の[`IntervalHeap<T>`]を構築する。
110    pub fn new() -> Self {
111        Self { data: vec![] }
112    }
113
114    /// 最小値の参照を返す。
115    pub fn min(&self) -> Option<&T> {
116        self.data.get(self.min_index())
117    }
118
119    /// 最大値の参照を返す。
120    pub fn max(&self) -> Option<&T> {
121        self.data.get(self.max_index())
122    }
123
124    /// 値`item`を挿入する。
125    pub fn push(&mut self, item: T) {
126        self.data.push(item);
127        self.bottom_up(self.data.len() - 1);
128    }
129
130    /// 最小値を削除して返す。
131    pub fn pop_min(&mut self) -> Option<T> {
132        if self.data.is_empty() {
133            None
134        } else {
135            let x = self.data.swap_remove(self.min_index());
136            let k = self.top_down(1);
137            self.bottom_up(k);
138            Some(x)
139        }
140    }
141
142    /// 最大値を削除して返す。
143    pub fn pop_max(&mut self) -> Option<T> {
144        if self.data.is_empty() {
145            None
146        } else {
147            let x = self.data.swap_remove(self.max_index());
148            let k = self.top_down(0);
149            self.bottom_up(k);
150            Some(x)
151        }
152    }
153
154    /// 要素数が`0`ならば`true`を返す。
155    pub fn is_empty(&self) -> bool {
156        self.data.is_empty()
157    }
158
159    /// 要素数を返す。
160    pub fn len(&self) -> usize {
161        self.data.len()
162    }
163}