Merge #340

340: feat: add combinations_with_replacement r=jswrenn a=tommilligan

As part of a personal project, I wanted to generate combinations of a set `n` of length `k`, but where elements may be repeated (i.e. may be present more than once in the output). I was not able to find an implementation of this in Rust anywhere - please correct me if I'm wrong!

The `CombinationsWithReplacement` iterator consumes a given `Iterator` of `Clone`-able items, and yields combinations of type `Vec<I::Item>` (the same api as `Combinations`). Suggestions for a more concise name welcome.

The algorithm is based on [this Stackoverflow question](https://stackoverflow.com/questions/127704/algorithm-to-return-all-combinations-of-k-elements-from-n) and [sample C code](https://stackoverflow.com/questions/561/how-to-use-combinations-of-sets-as-test-data#794)

Co-authored-by: Tom Milligan <[email protected]>

Author: bors[bot]

benches/combinations_with_replacement.rs

@@ -0,0 +1,34 @@
+#![feature(test)]
+
+extern crate itertools;
+extern crate test;
+
+use itertools::Itertools;
+use test::{black_box, Bencher};
+
+#[bench]
+fn comb_replacement_n10_k5(b: &mut Bencher) {
+    b.iter(|| {
+        for i in (0..10).combinations_with_replacement(5) {
+            black_box(i);
+        }
+    });
+}
+
+#[bench]
+fn comb_replacement_n5_k10(b: &mut Bencher) {
+    b.iter(|| {
+        for i in (0..5).combinations_with_replacement(10) {
+            black_box(i);
+        }
+    });
+}
+
+#[bench]
+fn comb_replacement_n10_k10(b: &mut Bencher) {
+    b.iter(|| {
+        for i in (0..10).combinations_with_replacement(10) {
+            black_box(i);
+        }
+    });
+}

src/combinations.rs

@@ -1,7 +1,7 @@
-
-use std::ops::Index;
 use std::fmt;
 
+use super::lazy_buffer::LazyBuffer;
+
 /// An iterator to iterate through all the `n`-length combinations in an iterator.
 ///
 /// See [`.combinations()`](../trait.Itertools.html#method.combinations) for more information.
@@ -98,68 +98,3 @@ impl<I> Iterator for Combinations<I>
         Some(result)
     }
 }
-
-#[derive(Debug)]
-struct LazyBuffer<I: Iterator> {
-    it: I,
-    done: bool,
-    buffer: Vec<I::Item>,
-}
-
-impl<I> LazyBuffer<I>
-    where I: Iterator
-{
-    pub fn new(it: I) -> LazyBuffer<I> {
-        let mut it = it;
-        let mut buffer = Vec::new();
-        let done;
-        if let Some(first) = it.next() {
-            buffer.push(first);
-            done = false;
-        } else {
-            done = true;
-        }
-        LazyBuffer {
-            it: it,
-            done: done,
-            buffer: buffer,
-        }
-    }
-
-    pub fn len(&self) -> usize {
-        self.buffer.len()
-    }
-
-    pub fn is_done(&self) -> bool {
-        self.done
-    }
-
-    pub fn get_next(&mut self) -> bool {
-        if self.done {
-            return false;
-        }
-        let next_item = self.it.next();
-        match next_item {
-            Some(x) => {
-                self.buffer.push(x);
-                true
-            }
-            None => {
-                self.done = true;
-                false
-            }
-        }
-    }
-}
-
-impl<I> Index<usize> for LazyBuffer<I>
-    where I: Iterator,
-          I::Item: Sized
-{
-    type Output = I::Item;
-
-    fn index<'b>(&'b self, _index: usize) -> &'b I::Item {
-        self.buffer.index(_index)
-    }
-}
-

src/combinations_with_replacement.rs

@@ -0,0 +1,108 @@
+use std::fmt;
+
+use super::lazy_buffer::LazyBuffer;
+
+/// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement.
+///
+/// See [`.combinations_with_replacement()`](../trait.Itertools.html#method.combinations_with_replacement) for more information.
+#[derive(Clone)]
+pub struct CombinationsWithReplacement<I>
+where
+    I: Iterator,
+    I::Item: Clone,
+{
+    n: usize,
+    indices: Vec<usize>,
+    // The current known max index value. This increases as pool grows.
+    max_index: usize,
+    pool: LazyBuffer<I>,
+    first: bool,
+}
+
+impl<I> fmt::Debug for CombinationsWithReplacement<I>
+where
+    I: Iterator + fmt::Debug,
+    I::Item: fmt::Debug + Clone,
+{
+    debug_fmt_fields!(Combinations, n, indices, max_index, pool, first);
+}
+
+impl<I> CombinationsWithReplacement<I>
+where
+    I: Iterator,
+    I::Item: Clone,
+{
+    /// Map the current mask over the pool to get an output combination
+    fn current(&self) -> Vec<I::Item> {
+        self.indices.iter().map(|i| self.pool[*i].clone()).collect()
+    }
+}
+
+/// Create a new `CombinationsWithReplacement` from a clonable iterator.
+pub fn combinations_with_replacement<I>(iter: I, n: usize) -> CombinationsWithReplacement<I>
+where
+    I: Iterator,
+    I::Item: Clone,
+{
+    let indices: Vec<usize> = vec![0; n];
+    let pool: LazyBuffer<I> = LazyBuffer::new(iter);
+
+    CombinationsWithReplacement {
+        n,
+        indices,
+        max_index: 0,
+        pool: pool,
+        first: true,
+    }
+}
+
+impl<I> Iterator for CombinationsWithReplacement<I>
+where
+    I: Iterator,
+    I::Item: Clone,
+{
+    type Item = Vec<I::Item>;
+    fn next(&mut self) -> Option<Self::Item> {
+        // If this is the first iteration, return early
+        if self.first {
+            // In empty edge cases, stop iterating immediately
+            return if self.n == 0 || self.pool.is_done() {
+                None
+            // Otherwise, yield the initial state
+            } else {
+                self.first = false;
+                Some(self.current())
+            };
+        }
+
+        // Check if we need to consume more from the iterator
+        // This will run while we increment our first index digit
+        if !self.pool.is_done() {
+            self.pool.get_next();
+            self.max_index = self.pool.len() - 1;
+        }
+
+        // Work out where we need to update our indices
+        let mut increment: Option<(usize, usize)> = None;
+        for (i, indices_int) in self.indices.iter().enumerate().rev() {
+            if indices_int < &self.max_index {
+                increment = Some((i, indices_int + 1));
+                break;
+            }
+        }
+
+        match increment {
+            // If we can update the indices further
+            Some((increment_from, increment_value)) => {
+                // We need to update the rightmost non-max value
+                // and all those to the right
+                for indices_index in increment_from..self.indices.len() {
+                    self.indices[indices_index] = increment_value
+                }
+                Some(self.current())
+            }
+            // Otherwise, we're done
+            None => None,
+        }
+    }
+}

src/lazy_buffer.rs

@@ -0,0 +1,67 @@
+use std::ops::Index;
+
+#[derive(Debug, Clone)]
+pub struct LazyBuffer<I: Iterator> {
+    it: I,
+    done: bool,
+    buffer: Vec<I::Item>,
+}
+
+impl<I> LazyBuffer<I>
+where
+    I: Iterator,
+{
+    pub fn new(it: I) -> LazyBuffer<I> {
+        let mut it = it;
+        let mut buffer = Vec::new();
+        let done;
+        if let Some(first) = it.next() {
+            buffer.push(first);
+            done = false;
+        } else {
+            done = true;
+        }
+        LazyBuffer {
+            it: it,
+            done: done,
+            buffer: buffer,
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.buffer.len()
+    }
+
+    pub fn is_done(&self) -> bool {
+        self.done
+    }
+
+    pub fn get_next(&mut self) -> bool {
+        if self.done {
+            return false;
+        }
+        let next_item = self.it.next();
+        match next_item {
+            Some(x) => {
+                self.buffer.push(x);
+                true
+            }
+            None => {
+                self.done = true;
+                false
+            }
+        }
+    }
+}
+
+impl<I> Index<usize> for LazyBuffer<I>
+where
+    I: Iterator,
+    I::Item: Sized,
+{
+    type Output = I::Item;
+
+    fn index<'b>(&'b self, _index: usize) -> &'b I::Item {
+        self.buffer.index(_index)
+    }
+}

src/lib.rs

@@ -102,6 +102,8 @@ pub mod structs {
     pub use adaptors::MultiProduct;
     #[cfg(feature = "use_std")]
     pub use combinations::Combinations;
+    #[cfg(feature = "use_std")]
+    pub use combinations_with_replacement::CombinationsWithReplacement;
     pub use cons_tuples_impl::ConsTuples;
     pub use exactly_one_err::ExactlyOneError;
     pub use format::{Format, FormatWith};
@@ -160,6 +162,8 @@ mod concat_impl;
 mod cons_tuples_impl;
 #[cfg(feature = "use_std")]
 mod combinations;
+#[cfg(feature = "use_std")]
+mod combinations_with_replacement;
 mod exactly_one_err;
 mod diff;
 mod format;
@@ -170,6 +174,8 @@ mod groupbylazy;
 mod intersperse;
 #[cfg(feature = "use_std")]
 mod kmerge_impl;
+#[cfg(feature = "use_std")]
+mod lazy_buffer;
 mod merge_join;
 mod minmax;
 #[cfg(feature = "use_std")]
@@ -1167,6 +1173,34 @@ pub trait Itertools : Iterator {
         combinations::combinations(self, n)
     }
 
+    /// Return an iterator that iterates over the `n`-length combinations of
+    /// the elements from an iterator, with replacement.
+    ///
+    /// Iterator element type is `Vec<Self::Item>`. The iterator produces a new Vec per iteration,
+    /// and clones the iterator elements.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let it = (1..4).combinations_with_replacement(2);
+    /// itertools::assert_equal(it, vec![
+    ///     vec![1, 1],
+    ///     vec![1, 2],
+    ///     vec![1, 3],
+    ///     vec![2, 2],
+    ///     vec![2, 3],
+    ///     vec![3, 3],
+    ///     ]);
+    /// ```
+    #[cfg(feature = "use_std")]
+    fn combinations_with_replacement(self, n: usize) -> CombinationsWithReplacement<Self>
+    where
+        Self: Sized,
+        Self::Item: Clone,
+    {
+        combinations_with_replacement::combinations_with_replacement(self, n)
+    }
+
     /// Return an iterator adaptor that pads the sequence to a minimum length of
     /// `min` by filling missing elements using a closure `f`.
     ///