diff --git a/src/extensions/collectible.rs b/src/extensions/collectible.rs index bf53cc5..ec29707 100644 --- a/src/extensions/collectible.rs +++ b/src/extensions/collectible.rs @@ -16,7 +16,6 @@ use std::iter::{Product, Sum}; pub trait Collectible: IntoIterator { // FIXME - implement these methods // cartesian_product - // combinations_repetitive /// Original collection type type This; @@ -155,6 +154,35 @@ pub trait Collectible: IntoIterator { Item: Clone, Self: Sized; + /// Creates a collection containing all combinations with repetition of specified size + /// from the elements of the original collection. + /// + /// The order or combined values is preserved for ordered collections. + /// + /// # Example + /// + /// ``` + /// use cantrip::*; + /// + /// let a = vec![1, 2, 3]; + /// + /// assert_eq!(a.combinations_repetitive(1), vec![vec![1], vec![2], vec![3]]); + /// assert_eq!( + /// a.combinations_repetitive(2), + /// vec![vec![1, 1], vec![1, 2], vec![1, 3], vec![2, 2], vec![2, 3], vec![3, 3]] + /// ); + /// assert_eq!( + /// a.combinations_repetitive(3), vec![ + /// vec![1, 1, 1], vec![1, 1, 2], vec![1, 1, 3], vec![1, 2, 2], vec![1, 2, 3], + /// vec![1, 3, 3], vec![2, 2, 2], vec![2, 2, 3], vec![2, 3, 3], vec![3, 3, 3], + /// ] + /// ); + /// ``` + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized; + /// Creates a collection containing an element /// specified number of times. /// @@ -1347,20 +1375,57 @@ where compute_combinations(&values, k) } +pub(crate) fn combinations_repetitive<'a, Item, Collection>( + iterator: impl Iterator, k: usize, +) -> Vec +where + Item: Clone + 'a, + Collection: FromIterator + Sized, +{ + if k == 0 { + return Vec::from_iter(iter::once(Collection::from_iter(iter::empty()))); + } + let values = Vec::from_iter(iterator); + let size = values.len(); + let mut combination = Vec::fill(0, k); + unfold(k > size, |done| { + if *done { + return None; + } + let result = Some(Collection::from_iter(combination.iter().map(|index| values[*index].clone()))); + let mut current_slot = k - 1; + while combination[current_slot] >= size - 1 { + if current_slot > 0 { + current_slot -= 1; + } else { + *done = true; + return result; + } + } + let current_index = combination[current_slot] + 1; + #[allow(clippy::needless_range_loop)] + for slot in current_slot..k { + combination[slot] = current_index; + } + result + }) + .collect() +} + pub(crate) fn compute_combinations<'a, Item, Collection>(values: &[&Item], k: usize) -> Vec where Item: Clone + 'a, Collection: FromIterator + Sized, { let size = values.len(); - let mut combination_indices = Vec::from_iter(0..k); + let mut combination = Vec::from_iter(0..k); unfold(k > size, |done| { if *done { return None; } - let result = Some(Collection::from_iter(combination_indices.iter().map(|index| values[*index].clone()))); + let result = Some(Collection::from_iter(combination.iter().map(|index| values[*index].clone()))); let mut current_slot = k - 1; - while combination_indices[current_slot] >= size + current_slot - k { + while combination[current_slot] >= size + current_slot - k { if current_slot > 0 { current_slot -= 1; } else { @@ -1368,11 +1433,11 @@ where return result; } } - let mut current_index = combination_indices[current_slot]; + let mut current_index = combination[current_slot]; #[allow(clippy::needless_range_loop)] for slot in current_slot..k { current_index += 1; - combination_indices[slot] = current_index; + combination[slot] = current_index; } result }) diff --git a/src/extensions/collections/binary_heap.rs b/src/extensions/collections/binary_heap.rs index ea0c2b0..948466f 100644 --- a/src/extensions/collections/binary_heap.rs +++ b/src/extensions/collections/binary_heap.rs @@ -88,6 +88,15 @@ impl Collectible for BinaryHeap { combinations(self.iter(), k) } + #[inline] + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized + { + combinations_repetitive(self.iter(), k) + } + #[inline] fn filter_map(&self, function: impl FnMut(&Item) -> Option) -> Self::This where diff --git a/src/extensions/collections/btree_set.rs b/src/extensions/collections/btree_set.rs index c0867c2..313dad6 100644 --- a/src/extensions/collections/btree_set.rs +++ b/src/extensions/collections/btree_set.rs @@ -88,6 +88,15 @@ impl Collectible for BTreeSet { combinations(self.iter(), k) } + #[inline] + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized + { + combinations_repetitive(self.iter(), k) + } + #[inline] fn filter_map(&self, function: impl FnMut(&Item) -> Option) -> Self::This where diff --git a/src/extensions/collections/hash_set.rs b/src/extensions/collections/hash_set.rs index 162b911..cea9dc1 100644 --- a/src/extensions/collections/hash_set.rs +++ b/src/extensions/collections/hash_set.rs @@ -89,6 +89,15 @@ impl Collectible for HashSet { combinations(self.iter(), k) } + #[inline] + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized + { + combinations_repetitive(self.iter(), k) + } + #[inline] fn filter_map(&self, function: impl FnMut(&Item) -> Option) -> Self::This where diff --git a/src/extensions/collections/linked_list.rs b/src/extensions/collections/linked_list.rs index cb8efc6..b2a7dac 100644 --- a/src/extensions/collections/linked_list.rs +++ b/src/extensions/collections/linked_list.rs @@ -101,6 +101,15 @@ impl Collectible for LinkedList { combinations(self.iter(), k) } + #[inline] + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized + { + combinations_repetitive(self.iter(), k) + } + #[inline] fn filter_map(&self, function: impl FnMut(&Item) -> Option) -> Self::This where diff --git a/src/extensions/collections/vec.rs b/src/extensions/collections/vec.rs index c62d8bf..e0a174c 100644 --- a/src/extensions/collections/vec.rs +++ b/src/extensions/collections/vec.rs @@ -114,6 +114,15 @@ impl Collectible for Vec { combinations(self.iter(), k) } + #[inline] + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized, + { + combinations_repetitive(self.iter(), k) + } + #[inline] fn filter_map(&self, function: impl FnMut(&Item) -> Option) -> Self::This where diff --git a/src/extensions/collections/vec_deque.rs b/src/extensions/collections/vec_deque.rs index f3d7f0f..a528741 100644 --- a/src/extensions/collections/vec_deque.rs +++ b/src/extensions/collections/vec_deque.rs @@ -101,6 +101,15 @@ impl Collectible for VecDeque { combinations(self.iter(), k) } + #[inline] + fn combinations_repetitive(&self, k: usize) -> Vec + where + Item: Clone, + Self: Sized + { + combinations_repetitive(self.iter(), k) + } + #[inline] fn filter_map(&self, function: impl FnMut(&Item) -> Option) -> Self::This where