From 9918693037dce8aa4bb6f08741b6812923486c18 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Wed, 19 Jun 2024 11:26:03 +0200 Subject: Merging upstream version 1.76.0+dfsg1. Signed-off-by: Daniel Baumann --- vendor/itertools/src/combinations.rs | 105 ++++++++++++++++++++++++++++------- 1 file changed, 85 insertions(+), 20 deletions(-) (limited to 'vendor/itertools/src/combinations.rs') diff --git a/vendor/itertools/src/combinations.rs b/vendor/itertools/src/combinations.rs index 68a59c5e4..d8b5351ec 100644 --- a/vendor/itertools/src/combinations.rs +++ b/vendor/itertools/src/combinations.rs @@ -4,6 +4,8 @@ use std::iter::FusedIterator; use super::lazy_buffer::LazyBuffer; use alloc::vec::Vec; +use crate::adaptors::checked_binomial; + /// An iterator to iterate through all the `k`-length combinations in an iterator. /// /// See [`.combinations()`](crate::Itertools::combinations) for more information. @@ -15,29 +17,29 @@ pub struct Combinations { } impl Clone for Combinations - where I: Clone + Iterator, - I::Item: Clone, +where + I: Clone + Iterator, + I::Item: Clone, { clone_fields!(indices, pool, first); } impl fmt::Debug for Combinations - where I: Iterator + fmt::Debug, - I::Item: fmt::Debug, +where + I: Iterator + fmt::Debug, + I::Item: fmt::Debug, { debug_fmt_fields!(Combinations, indices, pool, first); } /// Create a new `Combinations` from a clonable iterator. pub fn combinations(iter: I, k: usize) -> Combinations - where I: Iterator +where + I: Iterator, { - let mut pool = LazyBuffer::new(iter); - pool.prefill(k); - Combinations { indices: (0..k).collect(), - pool, + pool: LazyBuffer::new(iter), first: true, } } @@ -45,16 +47,22 @@ pub fn combinations(iter: I, k: usize) -> Combinations impl Combinations { /// Returns the length of a combination produced by this iterator. #[inline] - pub fn k(&self) -> usize { self.indices.len() } + pub fn k(&self) -> usize { + self.indices.len() + } /// Returns the (current) length of the pool from which combination elements are /// selected. This value can change between invocations of [`next`](Combinations::next). #[inline] - pub fn n(&self) -> usize { self.pool.len() } + pub fn n(&self) -> usize { + self.pool.len() + } - /// Returns a reference to the source iterator. + /// Returns a reference to the source pool. #[inline] - pub(crate) fn src(&self) -> &I { &self.pool.it } + pub(crate) fn src(&self) -> &LazyBuffer { + &self.pool + } /// Resets this `Combinations` back to an initial state for combinations of length /// `k` over the same pool data source. If `k` is larger than the current length @@ -68,7 +76,6 @@ impl Combinations { for i in 0..k { self.indices[i] = i; } - } else { for i in 0..self.indices.len() { self.indices[i] = i; @@ -77,15 +84,27 @@ impl Combinations { self.pool.prefill(k); } } + + pub(crate) fn n_and_count(self) -> (usize, usize) { + let Self { + indices, + pool, + first, + } = self; + let n = pool.count(); + (n, remaining_for(n, first, &indices).unwrap()) + } } impl Iterator for Combinations - where I: Iterator, - I::Item: Clone +where + I: Iterator, + I::Item: Clone, { type Item = Vec; fn next(&mut self) -> Option { if self.first { + self.pool.prefill(self.k()); if self.k() > self.n() { return None; } @@ -112,7 +131,7 @@ impl Iterator for Combinations // Increment index, and reset the ones to its right self.indices[i] += 1; - for j in i+1..self.indices.len() { + for j in i + 1..self.indices.len() { self.indices[j] = self.indices[j - 1] + 1; } } @@ -120,9 +139,55 @@ impl Iterator for Combinations // Create result vector based on the indices Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect()) } + + fn size_hint(&self) -> (usize, Option) { + let (mut low, mut upp) = self.pool.size_hint(); + low = remaining_for(low, self.first, &self.indices).unwrap_or(usize::MAX); + upp = upp.and_then(|upp| remaining_for(upp, self.first, &self.indices)); + (low, upp) + } + + #[inline] + fn count(self) -> usize { + self.n_and_count().1 + } } impl FusedIterator for Combinations - where I: Iterator, - I::Item: Clone -{} +where + I: Iterator, + I::Item: Clone, +{ +} + +/// For a given size `n`, return the count of remaining combinations or None if it would overflow. +fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option { + let k = indices.len(); + if n < k { + Some(0) + } else if first { + checked_binomial(n, k) + } else { + // https://en.wikipedia.org/wiki/Combinatorial_number_system + // http://www.site.uottawa.ca/~lucia/courses/5165-09/GenCombObj.pdf + + // The combinations generated after the current one can be counted by counting as follows: + // - The subsequent combinations that differ in indices[0]: + // If subsequent combinations differ in indices[0], then their value for indices[0] + // must be at least 1 greater than the current indices[0]. + // As indices is strictly monotonically sorted, this means we can effectively choose k values + // from (n - 1 - indices[0]), leading to binomial(n - 1 - indices[0], k) possibilities. + // - The subsequent combinations with same indices[0], but differing indices[1]: + // Here we can choose k - 1 values from (n - 1 - indices[1]) values, + // leading to binomial(n - 1 - indices[1], k - 1) possibilities. + // - (...) + // - The subsequent combinations with same indices[0..=i], but differing indices[i]: + // Here we can choose k - i values from (n - 1 - indices[i]) values: binomial(n - 1 - indices[i], k - i). + // Since subsequent combinations can in any index, we must sum up the aforementioned binomial coefficients. + + // Below, `n0` resembles indices[i]. + indices.iter().enumerate().try_fold(0usize, |sum, (i, n0)| { + sum.checked_add(checked_binomial(n - 1 - *n0, k - i)?) + }) + } +} -- cgit v1.2.3