count: combinations[_with_replacement], powerset

rust-itertools · Aug 13, 2023 · a3c62c9 · a3c62c9
1 parent 5750bc3
commit a3c62c9
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 31 deletions.
diff --git a/src/combinations.rs b/src/combinations.rs
@@ -53,6 +53,12 @@ impl<I: Iterator> Combinations<I> {
     #[inline]
     pub fn n(&self) -> usize { self.pool.len() }
 
+    /// Fill the pool to get its length.
+    pub(crate) fn real_n(&mut self) -> usize {
+        while self.pool.get_next() {}
+        self.pool.len()
+    }
+
     /// Returns a reference to the source iterator.
     #[inline]
     pub(crate) fn src(&self) -> &I { &self.pool.it }
@@ -78,6 +84,20 @@ impl<I: Iterator> Combinations<I> {
             self.pool.prefill(k);
         }
     }
+
+    fn remaining_for(&self, n: usize) -> Option<usize> {
+        let k = self.k();
+        if self.first {
+            binomial(n, k)
+        } else {
+            self.indices
+                .iter()
+                .enumerate()
+                .fold(Some(0), |sum, (k0, n0)| {
+                    sum.and_then(|s| s.checked_add(binomial(n - 1 - *n0, k - k0)?))
+                })
+        }
+    }
 }
 
 impl<I> Iterator for Combinations<I>
@@ -123,19 +143,12 @@ impl<I> Iterator for Combinations<I>
     }
 
     fn size_hint(&self) -> SizeHint {
-        let k = self.k();
-        size_hint::try_map(self.pool.size_hint(), |n| {
-            if self.first {
-                binomial(n, k)
-            } else {
-                self.indices
-                    .iter()
-                    .enumerate()
-                    .fold(Some(0), |sum, (k0, n0)| {
-                        sum.and_then(|s| s.checked_add(binomial(n - 1 - *n0, k - k0)?))
-                    })
-            }
-        })
+        size_hint::try_map(self.pool.size_hint(), |n| self.remaining_for(n))
+    }
+
+    fn count(mut self) -> usize {
+        let n = self.real_n();
+        self.remaining_for(n).expect("Iterator count greater than usize::MAX")
     }
 }
 

diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs
@@ -38,6 +38,21 @@ where
     fn current(&self) -> Vec<I::Item> {
         self.indices.iter().map(|i| self.pool[*i].clone()).collect()
     }
+
+    fn remaining_for(&self, n: usize) -> Option<usize> {
+        let k_perms = |n: usize, k: usize| binomial((n + k).saturating_sub(1), k);
+        let k = self.indices.len();
+        if self.first {
+            k_perms(n, k)
+        } else {
+            self.indices
+                .iter()
+                .enumerate()
+                .fold(Some(0), |sum, (k0, n0)| {
+                    sum.and_then(|s| s.checked_add(k_perms(n - 1 - *n0, k - k0)?))
+                })
+        }
+    }
 }
 
 /// Create a new `CombinationsWithReplacement` from a clonable iterator.
@@ -104,20 +119,13 @@ where
     }
 
     fn size_hint(&self) -> SizeHint {
-        let k_perms = |n: usize, k: usize| binomial((n + k).saturating_sub(1), k);
-        let k = self.indices.len();
-        size_hint::try_map(self.pool.size_hint(), |n| {
-            if self.first {
-                k_perms(n, k)
-            } else {
-                self.indices
-                    .iter()
-                    .enumerate()
-                    .fold(Some(0), |sum, (k0, n0)| {
-                        sum.and_then(|s| s.checked_add(k_perms(n - 1 - *n0, k - k0)?))
-                    })
-            }
-        })
+        size_hint::try_map(self.pool.size_hint(), |n| self.remaining_for(n))
+    }
+
+    fn count(mut self) -> usize {
+        while self.pool.get_next() {}
+        let n = self.pool.len();
+        self.remaining_for(n).expect("Iterator count greater than usize::MAX")
     }
 }
 

diff --git a/src/powerset.rs b/src/powerset.rs
@@ -3,7 +3,7 @@ use std::iter::FusedIterator;
 use std::usize;
 use alloc::vec::Vec;
 
-use super::combinations::{Combinations, combinations};
+use super::combinations::{Combinations, binomial, combinations};
 use super::size_hint;
 
 /// An iterator to iterate through the powerset of the elements from an iterator.
@@ -81,6 +81,13 @@ impl<I> Iterator for Powerset<I>
             (0, self_total.1)
         }
     }
+
+    fn count(mut self) -> usize {
+        let k = self.combs.k();
+        let n = self.combs.real_n();
+        // It could be `(1 << n) - self.pos` but `1 << n` might overflow.
+        self.combs.count() + (k + 1..=n).map(|k| binomial(n, k).unwrap()).sum::<usize>()
+    }
 }
 
 impl<I> FusedIterator for Powerset<I>

diff --git a/tests/test_std.rs b/tests/test_std.rs
@@ -910,13 +910,14 @@ fn combinations_zero() {
 }
 
 #[test]
-fn combinations_range_size_hint() {
+fn combinations_range_count() {
     for n in 0..6 {
         for k in 0..=n {
             let len = (n - k + 1..=n).product::<usize>() / (1..=k).product::<usize>();
             let mut it = (0..n).combinations(k);
             for count in (0..=len).rev() {
                 assert_eq!(it.size_hint(), (count, Some(count)));
+                assert_eq!(it.clone().count(), count);
                 assert_eq!(it.next().is_none(), count == 0);
             }
         }
@@ -930,13 +931,14 @@ fn permutations_zero() {
 }
 
 #[test]
-fn permutations_range_size_hint() {
+fn permutations_range_count() {
     for n in 0..6 {
         for k in 0..=n {
             let len: usize = (n - k + 1..=n).product();
             let mut it = (0..n).permutations(k);
             for count in (0..=len).rev() {
                 assert_eq!(it.size_hint(), (count, Some(count)));
+                assert_eq!(it.clone().count(), count);
                 assert_eq!(it.next().is_none(), count == 0);
             }
         }
@@ -977,13 +979,14 @@ fn combinations_with_replacement() {
 }
 
 #[test]
-fn combinations_with_replacement_range_size_hint() {
+fn combinations_with_replacement_range_count() {
     for n in 0..6 {
         for k in 0..=n {
             let len = (n..n + k).product::<usize>() / (1..=k).product::<usize>();
             let mut it = (0..n).combinations_with_replacement(k);
             for count in (0..=len).rev() {
                 assert_eq!(it.size_hint(), (count, Some(count)));
+                assert_eq!(it.clone().count(), count);
                 assert_eq!(it.next().is_none(), count == 0);
             }
         }
@@ -1005,6 +1008,11 @@ fn powerset() {
     assert_eq!((0..4).powerset().count(), 1 << 4);
     assert_eq!((0..8).powerset().count(), 1 << 8);
     assert_eq!((0..16).powerset().count(), 1 << 16);
+    let mut it = (0..8).powerset();
+    for count in (0..=(1 << 8)).rev() {
+        assert_eq!(it.clone().count(), count);
+        assert_eq!(it.next().is_none(), count == 0);
+    }
 }
 
 #[test]