diff --git a/src/liballoc/collections/binary_heap.rs b/src/liballoc/collections/binary_heap.rs index 03c9164fb9095..8e170d970bc57 100644 --- a/src/liballoc/collections/binary_heap.rs +++ b/src/liballoc/collections/binary_heap.rs @@ -665,6 +665,34 @@ impl BinaryHeap { pub fn drain_sorted(&mut self) -> DrainSorted<'_, T> { DrainSorted { inner: self } } + + /// Retains only the elements specified by the predicate. + /// + /// In other words, remove all elements `e` such that `f(&e)` returns + /// `false`. The elements are visited in unsorted (and unspecified) order. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(binary_heap_retain)] + /// use std::collections::BinaryHeap; + /// + /// let mut heap = BinaryHeap::from(vec![-10, -5, 1, 2, 4, 13]); + /// + /// heap.retain(|x| x % 2 == 0); // only keep even numbers + /// + /// assert_eq!(heap.into_sorted_vec(), [-10, 2, 4]) + /// ``` + #[unstable(feature = "binary_heap_retain", issue = "71503")] + pub fn retain(&mut self, f: F) + where + F: FnMut(&T) -> bool, + { + self.data.retain(f); + self.rebuild(); + } } impl BinaryHeap { diff --git a/src/liballoc/tests/binary_heap.rs b/src/liballoc/tests/binary_heap.rs index 057afd41824e7..62084ccf53c59 100644 --- a/src/liballoc/tests/binary_heap.rs +++ b/src/liballoc/tests/binary_heap.rs @@ -372,6 +372,14 @@ fn assert_covariance() { } } +#[test] +fn test_retain() { + let mut a = BinaryHeap::from(vec![-10, -5, 1, 2, 4, 13]); + a.retain(|x| x % 2 == 0); + + assert_eq!(a.into_sorted_vec(), [-10, 2, 4]) +} + // old binaryheap failed this test // // Integrity means that all elements are present after a comparison panics, diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs index ad6feaeebc67f..78d49558262e3 100644 --- a/src/liballoc/tests/lib.rs +++ b/src/liballoc/tests/lib.rs @@ -14,6 +14,7 @@ #![feature(binary_heap_drain_sorted)] #![feature(vec_remove_item)] #![feature(split_inclusive)] +#![feature(binary_heap_retain)] use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher};