1
use std::cmp::{max, Ordering};
2

            
3
/// An implicit binary heap that can efficiently find the min or max of its contained elements.
4
/// - Find min / find max: `O(1)`
5
/// - Pop min / pop max: `O(log(n))`
6
/// - Push: `O(log(n))`
7
/// - Heapify: `O(n)`
8
/// The comparison function can be changed after creation by just calling [`MmHeap::ify_by`] again.
9
/// Rust's standard [`std::collections::BinaryHeap`] claims an `O(1)` average complexity for push,
10
/// but as it notes in the documentation this is amortized over a randomly ordered sequence.
11
/// This heap is also a binary heap so for a randomly ordered sequence it could also be considered `O(1)`,
12
/// but for this, the standard binary heap, and indeed any binary heap, the worst case insertion time
13
/// is `O(logn)`.  Also, when a min (or minmax) heap is used to find the `m` largest elements in a large sequence of `n > m`
14
/// elements or visa versa, the expected cost of insertion drops from `1 - m/2^m` to ... even lower.
15
pub struct MmHeap<T> {
16
    buf: Vec<T>
17
}
18

            
19

            
20

            
21
impl<T> MmHeap<T> {
22
	/// Create an empty MmHeap
23
1000
    pub fn new() -> Self {
24
1000
        Self{buf: Vec::new()}
25
1000
    }
26

            
27
	/// Create an MmHeap out of a vector and immediately heapify it according to cmp
28
	/// Note that cmp takes references, which means the common case of storing references
29
	/// in the heap can require wrapping the comparison function since it will have to take
30
	/// references to references
31
    pub fn make(buf: Vec<T>, cmp: &impl Fn(&T, &T) -> Ordering) -> Self {
32
        let mut res = Self{buf};
33
        res.ify_by(cmp);
34
        res
35
    }
36

            
37
	/// Reorder the heap according to a new comparison function.
38
	/// It's impossible to add to the heap's private buffer directly, so this is only necessary
39
	/// when changing the comparison function.  When calling [`Self::push_by`] etc with the same comparison
40
	/// function, the heap will maintain its invariant
41
    pub fn ify_by(&mut self, cmp: &impl Fn(&T, &T) -> Ordering) {
42
        let nonleaf_idx_upper_bound = (usize::MAX >> 1) >> self.buf.len().leading_zeros();
43
        for i in (0..nonleaf_idx_upper_bound).rev()
44
            { self.sift_down_by(i, cmp) }
45
    }
46

            
47
	/// Get the number of elements in the heap
48
1000000
	pub fn len(&self) -> usize {
49
1000000
		self.buf.len()
50
1000000
	}
51

            
52
	/// Get the minimum element without removing it
53
    pub fn peek_min(&self) -> Option<&T> {
54
        self.buf.first()
55
    }
56

            
57
	/// Get the maximum element without removing it
58
    pub fn peek_max_by(&self, cmp: &impl Fn(&T, &T) -> Ordering) -> Option<&T> {
59
        match self.buf.get(1..3) {
60
            Some(slice_ref) => slice_ref.into_iter().max_by(|a, b|cmp(a, b)),
61
            None => self.buf.get(max(self.buf.len(), 1) - 1)
62
        }
63
    }
64

            
65
1000000
    fn sift_up_by(&mut self, mut i: usize, cmp: &impl Fn(&T, &T) -> Ordering) {
66
1000000
        if i == 0 || i >= self.buf.len()
67
999000
            { return }
68
        // nodes with index i will be in layer n where 2^n is the maximal power of 2 <= i + 1,
69
        // so we can check if n is odd/even by checking if the number of leading zeros in (i + 1)
70
        // is odd/even.  In odd layers, nodes should be >= their descendents, and in even layers, <=.
71
999000
        let mut ord = match (i + 1).leading_zeros()&1
72
979000
            { 1 => Ordering::Less, _ => Ordering::Greater };
73
999000
        let mut i1 = (i - 1) >> 1;
74
999000
        if cmp(&self.buf[i1], &self.buf[i]) == ord {
75
51024
            self.buf.swap(i, i1);
76
51024
            i = i1;
77
51024
            ord = ord.reverse()
78
947976
        }
79
2557106
        while i > 2 {
80
1790420
            i1 = (i - 3) >> 2;
81
1790420
            if cmp(&self.buf[i], &self.buf[i1]) == ord {
82
1558106
                self.buf.swap(i, i1);
83
1558106
                i = i1
84
232314
            } else { break }
85
        }
86
1000000
    }
87

            
88
950000
    fn sift_down_by(&mut self, mut i: usize, cmp: &impl Fn(&T, &T) -> Ordering) {
89
950000
        let ord = match (i + 1).leading_zeros()&1
90
950000
            { 1 => Ordering::Less, _ => Ordering::Greater };
91
2045126
        while 2*i + 1 < self.buf.len() {
92
            // Find m, the index of the extremal element among the children and grandchildren
93
            // of the element at index i. For min layers, extremal means
94
            // minimal, and for max layers it means maximal
95
1899476
            let mut m = 2*i + 1;
96
9363646
            for ii in [2*i + 2, 4*i + 3, 4*i + 4, 4*i + 5, 4*i + 6].into_iter().take_while(|&j|j<self.buf.len()) {
97
8545954
                if cmp(&self.buf[ii], &self.buf[m]) == ord
98
4817230
                    { m = ii }
99
            }
100
            // If m is a grandchild of i (as should be the case most of the time)
101
            // we may have to sift down farther after fixing up here
102
1899476
            if m > 2*i + 2 {
103
1854898
                if cmp(&self.buf[m], &self.buf[i]) == ord {
104
1095126
                    self.buf.swap(m, i);
105
1095126
                    let p = (m - 1) >> 1;
106
1095126
                    if cmp(&self.buf[p], &self.buf[m]) == ord
107
975980
                        { self.buf.swap(m, p) }
108
1095126
                    i = m;
109
759772
                } else { break }
110
            } else {// otherwise em is a direct child so it must be a leaf or its invariant would be wrong
111
44578
                if cmp(&self.buf[m], &self.buf[i]) == ord
112
36838
                    { self.buf.swap(m, i) }
113
44578
                break
114
            }
115
        }
116
950000
    }
117

            
118
	/// Insert an element into the heap
119
	/// Elements that compare equal are fine, but their order will be unspecified
120
1000000
    pub fn push_by(&mut self, e: T, cmp: &impl Fn(&T, &T) -> Ordering) {
121
1000000
        self.buf.push(e);
122
1000000
        self.sift_up_by(self.buf.len() - 1, cmp)
123
1000000
    }
124

            
125
	/// Get the minimal element and remove it
126
    pub fn pop_min_by(&mut self, cmp: &impl Fn(&T, &T) -> Ordering) -> Option<T> {
127
        self.pop_idx_by(0, cmp)
128
    }
129

            
130
	/// Get the maximal element and remove it
131
950000
    pub fn pop_max_by(&mut self, cmp: &impl Fn(&T, &T) -> Ordering) -> Option<T> {
132
950000
        match self.buf.get(1..3) {
133
950000
            Some(slice_ref) => self.pop_idx_by(1 + slice_ref.into_iter().enumerate().max_by(|(_i,a),(_j,b)|cmp(a,b)).unwrap().0, cmp),
134
            None => self.buf.pop()
135
        }
136
950000
    }
137

            
138
950000
    fn pop_idx_by(&mut self, i: usize, cmp: &impl Fn(&T, &T) -> Ordering) -> Option<T> {
139
950000
        let l = self.buf.len();
140
950000
        if i + 1 >= l
141
950000
            { return self.buf.pop() }
142
950000
        self.buf.swap(i, l - 1);
143
950000
        let res = self.buf.pop();
144
950000
        self.sift_down_by(i, cmp);
145
950000
        res
146
950000
    }
147

            
148
	/// Insert a new element and remove the min element in the resulting heap in a single operation
149
	/// Could be faster than [`Self::push_by`]` + [`Self::pop_min_by`] separately
150
	/// (the current implementation is only more efficient if the heap is empty, but it will always
151
	/// be at least as good)
152
    pub fn pushpop_min_by(&mut self, e: T, cmp: &impl Fn(&T, &T) -> Ordering) -> T {
153
        if self.buf.is_empty() { e } else {
154
            self.push_by(e, cmp);
155
            self.pop_min_by(cmp).unwrap()
156
        }
157
    }
158

            
159
	/// Insert a new element and remove the max element in the resulting heap in a single operation
160
950000
    pub fn pushpop_max_by(&mut self, e: T, cmp: &impl Fn(&T, &T) -> Ordering) -> T {
161
950000
        if self.buf.is_empty() { e } else {
162
950000
            self.push_by(e, cmp);
163
950000
            self.pop_max_by(cmp).unwrap()
164
        }
165
950000
    }
166

            
167
    /// Add all elements yielded from an iterator
168
    /// This is the same functionality provided by the [`Extend`] trait, except we need an additional
169
    /// argument (cmp) so we cannot implement that trait
170
500
    pub fn extend_by<U: IntoIterator<Item=T>>(&mut self, iter: U, cmp: &impl Fn(&T, &T) -> Ordering) {
171
500
        for x in iter {
172
            self.push_by(x, cmp)
173
        }
174
500
    }
175
}
176

            
177
impl<'a, T> IntoIterator for &'a MmHeap<T> {
178
    type Item = &'a T;
179
    type IntoIter = <&'a Vec<T> as IntoIterator>::IntoIter;
180
    fn into_iter(self) -> Self::IntoIter {
181
        (&self.buf).into_iter()
182
    }
183
}
184

            
185
impl<T> Into<Vec<T>> for MmHeap<T> {
186
1000
    fn into(self) -> Vec<T> {
187
1000
        self.buf
188
1000
    }
189
}
190