sim_lib_discrete_comb/
partition.rs

1//! Integer partitions of `n` in ascending-composition order (Kelleher's
2//! `accel_asc`). The iterator is genuinely lazy, so large `n` is safe to start.
3
4/// Iterator over the partitions of `n`, each as an ascending list of parts.
5#[derive(Debug, Clone)]
6pub struct IntegerPartitionIter {
7    a: Vec<usize>,
8    k: usize,
9    done: bool,
10    emit_empty: bool,
11}
12
13/// Construct a partition iterator for `n`. `n == 0` yields the empty partition.
14pub fn integer_partitions(n: usize) -> IntegerPartitionIter {
15    if n == 0 {
16        return IntegerPartitionIter {
17            a: Vec::new(),
18            k: 0,
19            done: true,
20            emit_empty: true,
21        };
22    }
23    let mut a = vec![0usize; n + 1];
24    a[1] = n;
25    IntegerPartitionIter {
26        a,
27        k: 1,
28        done: false,
29        emit_empty: false,
30    }
31}
32
33impl Iterator for IntegerPartitionIter {
34    type Item = Vec<usize>;
35
36    fn next(&mut self) -> Option<Self::Item> {
37        if self.emit_empty {
38            self.emit_empty = false;
39            return Some(Vec::new());
40        }
41        if self.done || self.k == 0 {
42            return None;
43        }
44        let mut k = self.k;
45        let x = self.a[k - 1] + 1;
46        let mut y = self.a[k] - 1;
47        k -= 1;
48        while x <= y {
49            self.a[k] = x;
50            y -= x;
51            k += 1;
52        }
53        self.a[k] = x + y;
54        let result = self.a[0..=k].to_vec();
55        self.k = k;
56        if k == 0 {
57            self.done = true;
58        }
59        Some(result)
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use crate::count::integer_partition_count;
67    use num_bigint::BigUint;
68
69    #[test]
70    fn count_matches_closed_form() {
71        let parts: Vec<_> = integer_partitions(5).collect();
72        assert_eq!(parts.len(), 7);
73        assert_eq!(
74            BigUint::from(parts.len() as u32),
75            integer_partition_count(5)
76        );
77        // Each partition sums to n.
78        for p in &parts {
79            assert_eq!(p.iter().sum::<usize>(), 5);
80        }
81    }
82
83    #[test]
84    fn first_and_last() {
85        let parts: Vec<_> = integer_partitions(5).collect();
86        assert_eq!(parts[0], vec![1, 1, 1, 1, 1]);
87        assert_eq!(parts[6], vec![5]);
88    }
89
90    #[test]
91    fn zero_yields_empty_partition() {
92        let parts: Vec<_> = integer_partitions(0).collect();
93        assert_eq!(parts, vec![Vec::<usize>::new()]);
94    }
95}
sim_lib_discrete_comb/partition.rs

sim_lib_discrete_comb/
partition.rs
