diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs index 8a36b2af76522..41df37ebd8d3a 100644 --- a/library/alloc/src/collections/binary_heap.rs +++ b/library/alloc/src/collections/binary_heap.rs @@ -143,6 +143,8 @@ #![allow(missing_docs)] #![stable(feature = "rust1", since = "1.0.0")] +use core::borrow::Borrow; +use core::cmp::Ordering; use core::fmt; use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen}; use core::mem::{self, swap, ManuallyDrop}; @@ -708,6 +710,96 @@ impl BinaryHeap { self.data.retain(f); self.rebuild(); } + + /// Returns `true` if the `BinaryHeap` contains the specified item. + /// + /// The item may be any borrowed form of the heap's item type, but + /// `PartialEq` on the borrowed form *must* match that for the item type. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_contains_remove)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap = BinaryHeap::new(); + /// heap.push("dog".to_string()); + /// heap.push("cat".to_string()); + /// heap.push("fish".to_string()); + /// + /// assert!(heap.contains("cat")); + /// assert!(!heap.contains("horse")); + /// ``` + /// + /// # Time complexity + /// + /// The worst case time complexity is *O*(*n*). + #[unstable(feature = "binary_heap_contains_remove", issue = "82001")] + pub fn contains(&self, item: &Q) -> bool + where + T: Borrow, + Q: PartialEq, + { + self.data.iter().any(|x| x.borrow() == item) + } + + /// Removes an item from the `BinaryHeap` and returns it, if it exists. + /// + /// The item may be any borrowed form of the heap's item type, but + /// `PartialEq` on the borrowed form *must* match that for the item type. If + /// there are multiple matching elements, there are no guarantees which of + /// the matching elements is removed on each call to `remove()`. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_contains_remove)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap = BinaryHeap::from(vec![1, 3, 5, 7, 9]); + /// + /// assert_eq!(heap.remove(&7), Some(7)); + /// assert_eq!(heap.remove(&7), None); + /// ``` + /// + /// # Time complexity + /// + /// The worst case time complexity is *O*(*n*). + #[unstable(feature = "binary_heap_contains_remove", issue = "82001")] + pub fn remove(&mut self, item: &Q) -> Option + where + T: Borrow, + Q: PartialEq, + { + self.data.iter().position(|x| x.borrow() == item).map(|pos| { + let last = self.data.len() - 1; + let comparison = self.data[last].cmp(&self.data[pos]); + let removed_item = self.data.swap_remove(pos); + match comparison { + Ordering::Less => { + // If the swapped item is less than the removed item, sift + // it down the heap to its correct position. + assert!(pos < last); + self.sift_down(pos); + } + Ordering::Greater => { + // If the swapped item is greater than the removed item, + // sift it up the heap to its correct position. + assert!(pos < last); + self.sift_up(0, pos); + } + Ordering::Equal => { + // If the swapped item is equal to the removed item, it's + // already in the correct position. + } + } + removed_item + }) + } } impl BinaryHeap { diff --git a/library/alloc/tests/binary_heap.rs b/library/alloc/tests/binary_heap.rs index ce794a9a4afa2..907b639fbe90f 100644 --- a/library/alloc/tests/binary_heap.rs +++ b/library/alloc/tests/binary_heap.rs @@ -392,6 +392,42 @@ fn test_retain() { assert_eq!(a.into_sorted_vec(), [-10, 2, 4]) } +#[test] +fn test_contains_remove() { + let mut a = BinaryHeap::from(vec![5, 7, 13, -2, -2, 8, -9]); + + assert!(a.contains(&7)); + assert_eq!(a.remove(&7), Some(7)); + assert_eq!(a.remove(&7), None); + assert!(!a.contains(&7)); + + assert!(a.contains(&-2)); + assert_eq!(a.remove(&-2), Some(-2)); + assert_eq!(a.remove(&-2), Some(-2)); + assert_eq!(a.remove(&-2), None); + assert!(!a.contains(&-2)); + + assert!(a.contains(&13)); + assert_eq!(a.remove(&13), Some(13)); + assert_eq!(a.remove(&13), None); + assert!(!a.contains(&13)); + + assert_eq!(a.into_sorted_vec(), [-9, 5, 8]); +} + +#[test] +fn test_empty_remove() { + let mut a = BinaryHeap::::new(); + assert_eq!(a.remove(&5), None); +} + +#[test] +fn test_singleton_remove() { + let mut a = BinaryHeap::from(vec![5]); + assert_eq!(a.remove(&5), Some(5)); + assert!(a.is_empty()); +} + // old binaryheap failed this test // // Integrity means that all elements are present after a comparison panics, diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index dd98f806451d8..26167f6fb6afd 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -21,6 +21,7 @@ #![feature(vecdeque_binary_search)] #![feature(slice_group_by)] #![feature(vec_extend_from_within)] +#![feature(binary_heap_contains_remove)] use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher};