Skip to content

Commit 7903f8e

Browse files
author
Gianmarco Garrisi
committed
Implement retain, retain_mut
1 parent a5b8496 commit 7903f8e

5 files changed

Lines changed: 243 additions & 0 deletions

File tree

‎src/double_priority_queue/mod.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,39 @@ where
421421
I: Hash + Eq,
422422
H: BuildHasher,
423423
{
424+
/// Retains only the elements specified by the `predicate`.
425+
///
426+
/// In other words, remove all elements e for which `predicate(&i, &p)` returns `false`.
427+
/// The elements are visited in arbitrary order.
428+
pub fn retain<F>(&mut self, predicate: F)
429+
where
430+
F: FnMut(&I, &P) -> bool,
431+
{
432+
self.store.retain(predicate);
433+
self.heap_build();
434+
}
435+
436+
/// Retains only the elements specified by the `predicate`.
437+
///
438+
/// In other words, remove all elements e for which `predicate(&mut i, &mut p)` returns `false`.
439+
/// The elements are visited in arbitrary order.
440+
///
441+
/// The `predicate` receives mutable references to both the item and
442+
/// the priority.
443+
///
444+
/// It's a logical error to change the item in a way
445+
/// that changes the result of `Hash` or `Eq`.
446+
///
447+
/// The `predicate` can change the priority. If the element is retained,
448+
/// it will have the updated one.
449+
pub fn retain_mut<F>(&mut self, predicate: F)
450+
where
451+
F: FnMut(&mut I, &mut P) -> bool,
452+
{
453+
self.store.retain_mut(predicate);
454+
self.heap_build();
455+
}
456+
424457
/// Removes the item with the lowest priority from
425458
/// the priority queue if the predicate returns `true`.
426459
///

‎src/priority_queue/mod.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,39 @@ where
331331
I: Hash + Eq,
332332
H: BuildHasher,
333333
{
334+
/// Retains only the elements specified by the `predicate`.
335+
///
336+
/// In other words, remove all elements e for which `predicate(&i, &p)` returns `false`.
337+
/// The elements are visited in arbitrary order.
338+
pub fn retain<F>(&mut self, predicate: F)
339+
where
340+
F: FnMut(&I, &P) -> bool,
341+
{
342+
self.store.retain(predicate);
343+
self.heap_build();
344+
}
345+
346+
/// Retains only the elements specified by the `predicate`.
347+
///
348+
/// In other words, remove all elements e for which `predicate(&mut i, &mut p)` returns `false`.
349+
/// The elements are visited in arbitrary order.
350+
///
351+
/// The `predicate` receives mutable references to both the item and
352+
/// the priority.
353+
///
354+
/// It's a logical error to change the item in a way
355+
/// that changes the result of `Hash` or `Eq`.
356+
///
357+
/// The `predicate` can change the priority. If the element is retained,
358+
/// it will have the updated one.
359+
pub fn retain_mut<F>(&mut self, predicate: F)
360+
where
361+
F: FnMut(&mut I, &mut P) -> bool,
362+
{
363+
self.store.retain_mut(predicate);
364+
self.heap_build();
365+
}
366+
334367
/// Removes the item with the greatest priority from
335368
/// the priority queue if the `predicate` returns `true`.
336369
///

‎src/store.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,25 @@ where
318318
I: Hash + Eq,
319319
H: BuildHasher,
320320
{
321+
pub fn retain<F>(&mut self, mut predicate: F)
322+
where
323+
F: FnMut(&I, &P) -> bool,
324+
{
325+
self.retain_mut(|i, p| predicate(&*i, &*p));
326+
}
327+
328+
pub fn retain_mut<F>(&mut self, predicate: F)
329+
where
330+
F: FnMut(&mut I, &mut P) -> bool,
331+
{
332+
self.map.retain2(predicate);
333+
if self.map.len() != self.size {
334+
self.size = self.map.len();
335+
self.heap = (0..self.size).into_iter().map(|i| Index(i)).collect();
336+
self.qp = (0..self.size).into_iter().map(|p| Position(p)).collect();
337+
}
338+
}
339+
321340
/// If the predicate returns true for the element in position `position`,
322341
/// remove it and swap it with the last element keeping a consistent
323342
/// state.

‎tests/double_priority_queue.rs‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -794,6 +794,85 @@ mod doublepq_tests {
794794
*/
795795
}
796796

797+
#[test]
798+
fn retain() {
799+
#[derive(Hash, PartialEq, Eq, Debug)]
800+
struct Animal {
801+
name: String,
802+
can_fly: bool,
803+
can_swim: bool,
804+
}
805+
806+
impl Animal {
807+
pub fn new(name: String, can_fly: bool, can_swim: bool) -> Self {
808+
Animal {
809+
name,
810+
can_fly,
811+
can_swim,
812+
}
813+
}
814+
}
815+
816+
let mut pq = DoublePriorityQueue::new();
817+
pq.push(Animal::new("dog".to_string(), false, true), 1);
818+
pq.push(Animal::new("cat".to_string(), false, false), 2);
819+
pq.push(Animal::new("bird".to_string(), true, false), 7);
820+
pq.push(Animal::new("fish".to_string(), false, true), 4);
821+
pq.push(Animal::new("cow".to_string(), false, false), 3);
822+
823+
pq.retain(|i, _| i.can_swim);
824+
825+
assert_eq!(
826+
pq.pop_max(),
827+
Some((Animal::new("fish".to_string(), false, true), 4))
828+
);
829+
assert_eq!(
830+
pq.pop_max(),
831+
Some((Animal::new("dog".to_string(), false, true), 1))
832+
);
833+
}
834+
835+
#[test]
836+
fn retain_mut() {
837+
#[derive(Hash, PartialEq, Eq, Debug)]
838+
struct Animal {
839+
name: String,
840+
can_fly: bool,
841+
can_swim: bool,
842+
}
843+
844+
impl Animal {
845+
pub fn new(name: String, can_fly: bool, can_swim: bool) -> Self {
846+
Animal {
847+
name,
848+
can_fly,
849+
can_swim,
850+
}
851+
}
852+
}
853+
854+
let mut pq = DoublePriorityQueue::new();
855+
pq.push(Animal::new("dog".to_string(), false, true), 1);
856+
pq.push(Animal::new("cat".to_string(), false, false), 2);
857+
pq.push(Animal::new("bird".to_string(), true, false), 7);
858+
pq.push(Animal::new("fish".to_string(), false, true), 4);
859+
pq.push(Animal::new("cow".to_string(), false, false), 3);
860+
861+
pq.retain_mut(|i, p| {
862+
*p += 10;
863+
i.can_swim
864+
});
865+
866+
assert_eq!(
867+
pq.pop_max(),
868+
Some((Animal::new("fish".to_string(), false, true), 14))
869+
);
870+
assert_eq!(
871+
pq.pop_max(),
872+
Some((Animal::new("dog".to_string(), false, true), 11))
873+
);
874+
}
875+
797876
#[test]
798877
fn into_sorted_iter() {
799878
let mut pq = DoublePriorityQueue::new();

‎tests/priority_queue.rs‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -681,6 +681,85 @@ mod pqueue_tests {
681681
*/
682682
}
683683

684+
#[test]
685+
fn retain() {
686+
#[derive(Hash, PartialEq, Eq, Debug)]
687+
struct Animal {
688+
name: String,
689+
can_fly: bool,
690+
can_swim: bool,
691+
}
692+
693+
impl Animal {
694+
pub fn new(name: String, can_fly: bool, can_swim: bool) -> Self {
695+
Animal {
696+
name,
697+
can_fly,
698+
can_swim,
699+
}
700+
}
701+
}
702+
703+
let mut pq = PriorityQueue::new();
704+
pq.push(Animal::new("dog".to_string(), false, true), 1);
705+
pq.push(Animal::new("cat".to_string(), false, false), 2);
706+
pq.push(Animal::new("bird".to_string(), true, false), 7);
707+
pq.push(Animal::new("fish".to_string(), false, true), 4);
708+
pq.push(Animal::new("cow".to_string(), false, false), 3);
709+
710+
pq.retain(|i, _| i.can_swim);
711+
712+
assert_eq!(
713+
pq.pop(),
714+
Some((Animal::new("fish".to_string(), false, true), 4))
715+
);
716+
assert_eq!(
717+
pq.pop(),
718+
Some((Animal::new("dog".to_string(), false, true), 1))
719+
);
720+
}
721+
722+
#[test]
723+
fn retain_mut() {
724+
#[derive(Hash, PartialEq, Eq, Debug)]
725+
struct Animal {
726+
name: String,
727+
can_fly: bool,
728+
can_swim: bool,
729+
}
730+
731+
impl Animal {
732+
pub fn new(name: String, can_fly: bool, can_swim: bool) -> Self {
733+
Animal {
734+
name,
735+
can_fly,
736+
can_swim,
737+
}
738+
}
739+
}
740+
741+
let mut pq = PriorityQueue::new();
742+
pq.push(Animal::new("dog".to_string(), false, true), 1);
743+
pq.push(Animal::new("cat".to_string(), false, false), 2);
744+
pq.push(Animal::new("bird".to_string(), true, false), 7);
745+
pq.push(Animal::new("fish".to_string(), false, true), 4);
746+
pq.push(Animal::new("cow".to_string(), false, false), 3);
747+
748+
pq.retain_mut(|i, p| {
749+
*p += 10;
750+
i.can_swim
751+
});
752+
753+
assert_eq!(
754+
pq.pop(),
755+
Some((Animal::new("fish".to_string(), false, true), 14))
756+
);
757+
assert_eq!(
758+
pq.pop(),
759+
Some((Animal::new("dog".to_string(), false, true), 11))
760+
);
761+
}
762+
684763
#[test]
685764
fn into_sorted_iter() {
686765
let mut pq = PriorityQueue::new();

0 commit comments

Comments
 (0)