Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify push_pop_{max, min} and replace_{max, min} implementations #8

Merged
merged 3 commits into from
Dec 19, 2021
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 102 additions & 49 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ impl<T: Ord> MinMaxHeap<T> {
} else {
Some(PeekMinMut {
heap: self,
removed: false,
sift: false,
})
}
}
Expand All @@ -141,7 +141,7 @@ impl<T: Ord> MinMaxHeap<T> {
self.find_max().map(move |i| PeekMaxMut {
heap: self,
max_index: i,
removed: false,
sift: false,
})
}

Expand Down Expand Up @@ -212,12 +212,11 @@ impl<T: Ord> MinMaxHeap<T> {
///
/// *O*(log *n*).
pub fn push_pop_min(&mut self, mut element: T) -> T {
if self.is_empty() { return element; }

if element < self.0[0] { return element; }

mem::swap(&mut element, &mut self.0[0]);
self.trickle_down_min(0);
if let Some(mut min) = self.peek_min_mut() {
if element > *min {
mem::swap(&mut element, &mut min);
}
}
element
}

Expand Down Expand Up @@ -245,18 +244,11 @@ impl<T: Ord> MinMaxHeap<T> {
///
/// *O*(log *n*).
pub fn push_pop_max(&mut self, mut element: T) -> T {
if let Some(i) = self.find_max() {
if element > self.0[i] { return element }

mem::swap(&mut element, &mut self.0[i]);

if self.0[i] < self.0[0] {
self.0.swap(0, i);
if let Some(mut max) = self.peek_max_mut() {
if element < *max {
mem::swap(&mut element, &mut max);
}

self.trickle_down_max(i);
}

element
}

Expand Down Expand Up @@ -284,12 +276,15 @@ impl<T: Ord> MinMaxHeap<T> {
/// <struct.MinMaxHeap.html#method.push_pop_min>
///
/// *O*(log *n*).
pub fn replace_min(&mut self, mut element: T) -> Option<T> { if
self.is_empty() { self.push(element); return None; }
pub fn replace_min(&mut self, mut element: T) -> Option<T> {
if let Some(mut min) = self.peek_min_mut() {
mem::swap(&mut element, &mut min);
return Some(element);
}

mem::swap(&mut element, &mut self.0[0]);
self.trickle_down_min(0);
Some(element)
// Heap was empty, so no reordering is neessary
self.0.push(element);
None
}

/// Pops the maximum element and pushes a new element, in an
Expand Down Expand Up @@ -317,19 +312,22 @@ impl<T: Ord> MinMaxHeap<T> {
///
/// *O*(log *n*).
pub fn replace_max(&mut self, mut element: T) -> Option<T> {
if let Some(i) = self.find_max() {
mem::swap(&mut element, &mut self.0[i]);

if self.0[i] < self.0[0] {
self.0.swap(0, i);
if let Some(mut max) = self.peek_max_mut() {
// If `element` is the new min, swap it with the current min
// (unless the min is the same as the max)
if max.heap.len() > 1 {
let min = &mut max.heap.0[0];
if element < *min {
mem::swap(&mut element, min);
}
}

self.trickle_down_max(i);
Some(element)
} else {
self.push(element);
None
mem::swap(&mut element, &mut max);
return Some(element);
}

// Heap was empty, so no reordering is neessary
self.0.push(element);
None
}

/// Returns an ascending (sorted) vector, reusing the heap’s
Expand Down Expand Up @@ -480,7 +478,7 @@ impl<T> MinMaxHeap<T> {
/// Returns a draining iterator over the min-max-heap’s elements in
/// descending (max-first) order.
///
/// *O*(1) on creation, and *O*(log *n*) for each `next()` operation.
/// *<<1) on creation, and *O*(log *n*) for each `next()` operation.
pub fn drain_desc(&mut self) -> DrainDesc<T> {
DrainDesc(self)
}
Expand Down Expand Up @@ -690,7 +688,7 @@ impl<'a, T: Ord + Clone + 'a> Extend<&'a T> for MinMaxHeap<T> {
/// [`MinMaxHeap`]: struct.MinMaxHeap.html
pub struct PeekMinMut<'a, T: 'a + Ord> {
heap: &'a mut MinMaxHeap<T>,
removed: bool,
sift: bool,
}

impl<T: Ord + fmt::Debug> fmt::Debug for PeekMinMut<'_, T> {
Expand All @@ -703,7 +701,7 @@ impl<T: Ord + fmt::Debug> fmt::Debug for PeekMinMut<'_, T> {

impl<'a, T: Ord> Drop for PeekMinMut<'a, T> {
fn drop(&mut self) {
if !self.removed {
if self.sift {
self.heap.trickle_down_min(0);
}
}
Expand All @@ -721,6 +719,7 @@ impl<'a, T: Ord> Deref for PeekMinMut<'a, T> {
impl<'a, T: Ord> DerefMut for PeekMinMut<'a, T> {
fn deref_mut(&mut self) -> &mut T {
debug_assert!(!self.heap.is_empty());
self.sift = true;
// SAFE: PeekMinMut is only instantiated for non-empty heaps
unsafe { self.heap.0.get_unchecked_mut(0) }
}
Expand All @@ -729,9 +728,9 @@ impl<'a, T: Ord> DerefMut for PeekMinMut<'a, T> {
impl<'a, T: Ord> PeekMinMut<'a, T> {
/// Removes the peeked value from the heap and returns it.
pub fn pop(mut self) -> T {
let value = self.heap.pop_min().unwrap();
self.removed = true;
value
// Sift is unnecessary since pop_min() already reorders heap
self.sift = false;
self.heap.pop_min().unwrap()
}
}

Expand All @@ -746,7 +745,7 @@ impl<'a, T: Ord> PeekMinMut<'a, T> {
pub struct PeekMaxMut<'a, T: 'a + Ord> {
heap: &'a mut MinMaxHeap<T>,
max_index: usize,
removed: bool,
sift: bool,
}

impl<T: Ord + fmt::Debug> fmt::Debug for PeekMaxMut<'_, T> {
Expand All @@ -759,14 +758,18 @@ impl<T: Ord + fmt::Debug> fmt::Debug for PeekMaxMut<'_, T> {

impl<'a, T: Ord> Drop for PeekMaxMut<'a, T> {
fn drop(&mut self) {
if !self.removed {
if self.sift {
let mut hole = Hole::new(&mut self.heap.0, self.max_index);

if hole.element() < hole.get_parent() {
hole.swap_with_parent();
}
// If the hole doesn't have a parent, the heap has a single element,
// so no reordering is necessary
if hole.has_parent() {
if hole.element() < hole.get_parent() {
hole.swap_with_parent();
}

hole.trickle_down_max();
hole.trickle_down_max();
}
}
}
}
Expand All @@ -783,6 +786,7 @@ impl<'a, T: Ord> Deref for PeekMaxMut<'a, T> {
impl<'a, T: Ord> DerefMut for PeekMaxMut<'a, T> {
fn deref_mut(&mut self) -> &mut T {
debug_assert!(self.max_index < self.heap.len());
self.sift = true;
// SAFE: PeekMaxMut is only instantiated for non-empty heaps
unsafe { self.heap.0.get_unchecked_mut(self.max_index) }
}
Expand All @@ -791,9 +795,9 @@ impl<'a, T: Ord> DerefMut for PeekMaxMut<'a, T> {
impl<'a, T: Ord> PeekMaxMut<'a, T> {
/// Removes the peeked value from the heap and returns it.
pub fn pop(mut self) -> T {
let value = self.heap.pop_max().unwrap();
self.removed = true;
value
// Sift is unnecessary since pop_max() already reorders heap
self.sift = false;
self.heap.pop_max().unwrap()
}
}

Expand Down Expand Up @@ -900,6 +904,31 @@ mod tests {
result
}

#[test]
fn replace_min() {
let mut h = MinMaxHeap::from(vec![1, 2]);
assert_eq!(Some(1), h.replace_min(0));
assert_eq!(Some(&0), h.peek_min());
assert_eq!(Some(&2), h.peek_max());

assert_eq!(Some(0), h.replace_min(3));
assert_eq!(Some(&2), h.peek_min());
assert_eq!(Some(&3), h.peek_max());
}

#[test]
fn replace_min_edge_cases() {
let mut empty_heap = MinMaxHeap::new();
assert_eq!(None, empty_heap.replace_min(1));
assert_eq!(Some(1), empty_heap.pop_min());
assert_eq!(None, empty_heap.pop_min());

let mut one_element_heap = MinMaxHeap::from(vec![2]);
assert_eq!(Some(2), one_element_heap.replace_min(1));
assert_eq!(Some(1), one_element_heap.pop_min());
assert_eq!(None, one_element_heap.pop_min());
}

#[test]
fn replace_max() {
let mut h = MinMaxHeap::from(vec![1, 2]);
Expand All @@ -912,6 +941,19 @@ mod tests {
assert_eq!(Some(&1), h.peek_max());
}

#[test]
fn replace_max_edge_cases() {
let mut empty_heap = MinMaxHeap::new();
assert_eq!(None, empty_heap.replace_max(1));
assert_eq!(Some(1), empty_heap.pop_max());
assert_eq!(None, empty_heap.pop_max());

let mut one_element_heap = MinMaxHeap::from(vec![1]);
assert_eq!(Some(1), one_element_heap.replace_max(2));
assert_eq!(Some(2), one_element_heap.pop_max());
assert_eq!(None, one_element_heap.pop_max());
}

#[test]
fn peek_min_mut() {
let mut h = MinMaxHeap::from(vec![2, 3, 4]);
Expand Down Expand Up @@ -944,6 +986,17 @@ mod tests {
assert_eq!(Some(&0), h.peek_max());
}

#[test]
fn peek_max_mut_one() {
let mut h = MinMaxHeap::from(vec![1]);
{
let mut max = h.peek_max_mut().unwrap();
assert_eq!(*max, 1);
*max = 2;
}
assert_eq!(h.peek_max(), Some(&2));
}

#[test]
fn push_pop_max() {
let mut h = MinMaxHeap::from(vec![1, 2]);
Expand Down