diff --git a/src/lib.rs b/src/lib.rs index 834a48dea..6e29cb123 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -126,6 +126,8 @@ pub mod structs { #[cfg(feature = "use_alloc")] pub use crate::permutations::Permutations; #[cfg(feature = "use_alloc")] + pub use crate::multiset_permutations::MultisetPermutations; + #[cfg(feature = "use_alloc")] pub use crate::powerset::Powerset; pub use crate::process_results_impl::ProcessResults; #[cfg(feature = "use_alloc")] @@ -216,6 +218,8 @@ mod peeking_take_while; #[cfg(feature = "use_alloc")] mod permutations; #[cfg(feature = "use_alloc")] +mod multiset_permutations; +#[cfg(feature = "use_alloc")] mod powerset; mod process_results_impl; #[cfg(feature = "use_alloc")] @@ -1841,6 +1845,55 @@ pub trait Itertools: Iterator { permutations::permutations(self, k) } + /// Return an iterator adaptor that iterates over all unique multiset permutations + /// of the elements from an iterator. + /// + /// Iterator element type is `Vec` with equal length to the source iterator. + /// The iterator produces a new Vec per iteration, and clones the iterator elements. + /// + /// + /// + /// ``` + /// use itertools::Itertools; + /// + /// let mut it = vec![1, 4, 2, 1].into_iter().multiset_permutations(); + /// assert_eq!(it.next(), Some(vec![4, 2, 1, 1])); + /// assert_eq!(it.next(), Some(vec![1, 4, 2, 1])); + /// assert_eq!(it.next(), Some(vec![4, 1, 2, 1])); + /// assert_eq!(it.next(), Some(vec![1, 4, 1, 2])); + /// assert_eq!(it.next(), Some(vec![1, 1, 4, 2])); + /// assert_eq!(it.next(), Some(vec![4, 1, 1, 2])); + /// assert_eq!(it.next(), Some(vec![2, 4, 1, 1])); + /// assert_eq!(it.next(), Some(vec![1, 2, 4, 1])); + /// assert_eq!(it.next(), Some(vec![2, 1, 4, 1])); + /// assert_eq!(it.next(), Some(vec![1, 2, 1, 4])); + /// assert_eq!(it.next(), Some(vec![1, 1, 2, 4])); + /// assert_eq!(it.next(), Some(vec![2, 1, 1, 4])); + /// assert_eq!(it.next(), None); + /// ``` + /// + /// If the source iterator is empty, the resultant iterator adaptor + /// will only contain an empty vector. + /// + /// ``` + /// use itertools::Itertools; + /// use itertools::MultisetPermutations; + /// + /// let mut it: MultisetPermutations = vec![].into_iter().multiset_permutations(); + /// assert_eq!(it.next(), Some(vec![])); + /// assert_eq!(it.next(), None); + /// ``` + /// + /// Note: The source iterator is collected completely + #[cfg(feature = "use_alloc")] + fn multiset_permutations(self) -> MultisetPermutations + where + Self: Sized, + Self::Item: Ord, + { + multiset_permutations::multiset_permutations(self) + } + /// Return an iterator that iterates through the powerset of the elements from an /// iterator. /// diff --git a/src/multiset_permutations.rs b/src/multiset_permutations.rs new file mode 100644 index 000000000..1b4c0fa84 --- /dev/null +++ b/src/multiset_permutations.rs @@ -0,0 +1,150 @@ +use std::iter::FromIterator; + +/// An iterator adaptor that iterates through all the unique multiset permutations of the iterator. +/// The supplied iterator is fully consumed, so it must be finite. +/// +/// See [`.multiset_permutations()`](crate::Itertools::multiset_permutations) for +/// more information. +#[derive(Debug, Clone)] +pub struct MultisetPermutations { + buffer: Vec, + start: bool, + index: usize, +} + +pub fn multiset_permutations(iter: I) -> MultisetPermutations +where + I: Iterator, + I::Item: Ord, +{ + let mut buffer = Vec::from_iter(iter); + buffer.sort_unstable_by(|a, b| b.cmp(a)); + let length = buffer.len(); + MultisetPermutations { + buffer: buffer, + start: true, + index: length.saturating_sub(2), + } +} + +impl Iterator for MultisetPermutations +where + I: Ord, +{ + type Item = Vec; + + fn next(&mut self) -> Option { + // Start iteration with buffer itself + if self.start { + self.start = false; + return Some(self.buffer.clone()); + } + + // Exhausted iteration + let has_two_next = self.index + 2 < self.buffer.len(); + if !has_two_next + && (self.buffer.len() <= self.index + 1 + || self.buffer[0] <= self.buffer[self.index + 1]) + { + return None; + } + + // Determine shift index + let shift_index = if has_two_next && self.buffer[self.index + 2] <= self.buffer[self.index] + { + self.index + 2 + } else { + self.index + 1 + }; + + // Prefix shift + let shift_elem = self.buffer[shift_index]; + let mut swap_index = shift_index; + while swap_index > 0 { + self.buffer[swap_index] = self.buffer[swap_index - 1]; + swap_index -= 1; + } + self.buffer[0] = shift_elem; + + // Update index + if self.buffer[0] < self.buffer[1] { + self.index = 0; + } else { + self.index += 1; + } + + Some(self.buffer.clone()) + } +} + +#[cfg(test)] +mod tests { + use crate::Itertools; + + #[test] + fn test1() { + let mut iter = vec![1, 4, 2, 1].into_iter().multiset_permutations(); + assert_eq!(iter.next(), Some(vec![4, 2, 1, 1])); + assert_eq!(iter.next(), Some(vec![1, 4, 2, 1])); + assert_eq!(iter.next(), Some(vec![4, 1, 2, 1])); + assert_eq!(iter.next(), Some(vec![1, 4, 1, 2])); + assert_eq!(iter.next(), Some(vec![1, 1, 4, 2])); + assert_eq!(iter.next(), Some(vec![4, 1, 1, 2])); + assert_eq!(iter.next(), Some(vec![2, 4, 1, 1])); + assert_eq!(iter.next(), Some(vec![1, 2, 4, 1])); + assert_eq!(iter.next(), Some(vec![2, 1, 4, 1])); + assert_eq!(iter.next(), Some(vec![1, 2, 1, 4])); + assert_eq!(iter.next(), Some(vec![1, 1, 2, 4])); + assert_eq!(iter.next(), Some(vec![2, 1, 1, 4])); + assert_eq!(iter.next(), None); + } + + #[test] + fn test3() { + use rand::Rng; + let n = 12; + let mut rng = rand::thread_rng(); + let vec: Vec = (0..n).map(|_| rng.gen_range(0, 10)).collect(); + + let mut permutations1 = vec + .clone() + .into_iter() + .permutations(n) + .unique() + .collect_vec(); + let mut permutations2 = vec.into_iter().multiset_permutations().collect_vec(); + + permutations1.sort(); + permutations2.sort(); + assert_eq!(permutations1, permutations2); + } + + #[test] + fn test4() { + let mut iter = vec![0, 0, 1].into_iter().multiset_permutations(); + assert_eq!(iter.next(), Some(vec![1, 0, 0])); + assert_eq!(iter.next(), Some(vec![0, 1, 0])); + assert_eq!(iter.next(), Some(vec![0, 0, 1])); + assert_eq!(iter.next(), None); + } + + #[test] + fn test5() { + let mut iter = vec![1, 1].into_iter().multiset_permutations(); + assert_eq!(iter.next(), Some(vec![1, 1])); + assert_eq!(iter.next(), None); + } + + #[test] + fn test6() { + let iter = "MISSISSIPPI".chars().multiset_permutations(); + assert_eq!(iter.count(), 34650); // 34650 = 11! / (1! * 2! * 4! * 4!) + } + + #[test] + fn test7() { + let mut iter: crate::MultisetPermutations = vec![].into_iter().multiset_permutations(); + assert_eq!(iter.next(), Some(vec![])); + assert_eq!(iter.next(), None); + } +} diff --git a/tests/quick.rs b/tests/quick.rs index 672901e7c..c4041af33 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -740,6 +740,22 @@ quickcheck! { assert_eq!(expected, actual); } + + fn correct_multiset_permutations(a: usize, b: usize, c: usize, d: usize) -> () { + let n1 = a % 2; + let n2 = b % 4; + let n3 = c % 2; + let n4 = d % 5; + let mut vec = vec![0; n1]; + vec.append(&mut vec![1; n2]); + vec.append(&mut vec![2; n3]); + vec.append(&mut vec![3; n4]); + let n = vec.len(); + let unique_permutations = vec.iter().permutations(n).unique().sorted().collect_vec(); + let multiset_permutations = vec.iter().multiset_permutations().sorted().collect_vec(); + assert_eq!(unique_permutations, multiset_permutations); + } + } quickcheck! {