/rust/registry/src/index.crates.io-1949cf8c6b5b557f/itertools-0.14.0/src/powerset.rs
Line | Count | Source |
1 | | use alloc::vec::Vec; |
2 | | use std::fmt; |
3 | | use std::iter::FusedIterator; |
4 | | |
5 | | use super::combinations::{combinations, Combinations}; |
6 | | use crate::adaptors::checked_binomial; |
7 | | use crate::size_hint::{self, SizeHint}; |
8 | | |
9 | | /// An iterator to iterate through the powerset of the elements from an iterator. |
10 | | /// |
11 | | /// See [`.powerset()`](crate::Itertools::powerset) for more |
12 | | /// information. |
13 | | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] |
14 | | pub struct Powerset<I: Iterator> { |
15 | | combs: Combinations<I>, |
16 | | } |
17 | | |
18 | | impl<I> Clone for Powerset<I> |
19 | | where |
20 | | I: Clone + Iterator, |
21 | | I::Item: Clone, |
22 | | { |
23 | | clone_fields!(combs); |
24 | | } |
25 | | |
26 | | impl<I> fmt::Debug for Powerset<I> |
27 | | where |
28 | | I: Iterator + fmt::Debug, |
29 | | I::Item: fmt::Debug, |
30 | | { |
31 | | debug_fmt_fields!(Powerset, combs); |
32 | | } |
33 | | |
34 | | /// Create a new `Powerset` from a clonable iterator. |
35 | 0 | pub fn powerset<I>(src: I) -> Powerset<I> |
36 | 0 | where |
37 | 0 | I: Iterator, |
38 | 0 | I::Item: Clone, |
39 | | { |
40 | 0 | Powerset { |
41 | 0 | combs: combinations(src, 0), |
42 | 0 | } |
43 | 0 | } |
44 | | |
45 | | impl<I: Iterator> Powerset<I> { |
46 | | /// Returns true if `k` has been incremented, false otherwise. |
47 | 0 | fn increment_k(&mut self) -> bool { |
48 | 0 | if self.combs.k() < self.combs.n() || self.combs.k() == 0 { |
49 | 0 | self.combs.reset(self.combs.k() + 1); |
50 | 0 | true |
51 | | } else { |
52 | 0 | false |
53 | | } |
54 | 0 | } |
55 | | } |
56 | | |
57 | | impl<I> Iterator for Powerset<I> |
58 | | where |
59 | | I: Iterator, |
60 | | I::Item: Clone, |
61 | | { |
62 | | type Item = Vec<I::Item>; |
63 | | |
64 | 0 | fn next(&mut self) -> Option<Self::Item> { |
65 | 0 | if let Some(elt) = self.combs.next() { |
66 | 0 | Some(elt) |
67 | 0 | } else if self.increment_k() { |
68 | 0 | self.combs.next() |
69 | | } else { |
70 | 0 | None |
71 | | } |
72 | 0 | } |
73 | | |
74 | 0 | fn nth(&mut self, mut n: usize) -> Option<Self::Item> { |
75 | | loop { |
76 | 0 | match self.combs.try_nth(n) { |
77 | 0 | Ok(item) => return Some(item), |
78 | 0 | Err(steps) => { |
79 | 0 | if !self.increment_k() { |
80 | 0 | return None; |
81 | 0 | } |
82 | 0 | n -= steps; |
83 | | } |
84 | | } |
85 | | } |
86 | 0 | } |
87 | | |
88 | 0 | fn size_hint(&self) -> SizeHint { |
89 | 0 | let k = self.combs.k(); |
90 | | // Total bounds for source iterator. |
91 | 0 | let (n_min, n_max) = self.combs.src().size_hint(); |
92 | 0 | let low = remaining_for(n_min, k).unwrap_or(usize::MAX); |
93 | 0 | let upp = n_max.and_then(|n| remaining_for(n, k)); |
94 | 0 | size_hint::add(self.combs.size_hint(), (low, upp)) |
95 | 0 | } |
96 | | |
97 | 0 | fn count(self) -> usize { |
98 | 0 | let k = self.combs.k(); |
99 | 0 | let (n, combs_count) = self.combs.n_and_count(); |
100 | 0 | combs_count + remaining_for(n, k).unwrap() |
101 | 0 | } |
102 | | |
103 | 0 | fn fold<B, F>(self, mut init: B, mut f: F) -> B |
104 | 0 | where |
105 | 0 | F: FnMut(B, Self::Item) -> B, |
106 | | { |
107 | 0 | let mut it = self.combs; |
108 | 0 | if it.k() == 0 { |
109 | 0 | init = it.by_ref().fold(init, &mut f); |
110 | 0 | it.reset(1); |
111 | 0 | } |
112 | 0 | init = it.by_ref().fold(init, &mut f); |
113 | | // n is now known for sure because k >= 1 and all k-combinations have been generated. |
114 | 0 | for k in it.k() + 1..=it.n() { |
115 | 0 | it.reset(k); |
116 | 0 | init = it.by_ref().fold(init, &mut f); |
117 | 0 | } |
118 | 0 | init |
119 | 0 | } |
120 | | } |
121 | | |
122 | | impl<I> FusedIterator for Powerset<I> |
123 | | where |
124 | | I: Iterator, |
125 | | I::Item: Clone, |
126 | | { |
127 | | } |
128 | | |
129 | 0 | fn remaining_for(n: usize, k: usize) -> Option<usize> { |
130 | 0 | (k + 1..=n).try_fold(0usize, |sum, i| sum.checked_add(checked_binomial(n, i)?)) |
131 | 0 | } |