sim_lib_discrete_comb/
combination.rs

1//! `k`-combinations of `{0, ..., n-1}` in lexicographic order, with rank/unrank.
2//!
3//! A combination is a strictly ascending list of `k` indices. Rank and unrank
4//! agree with the iterator's lexicographic order.
5
6use crate::count::binomial;
7use crate::error::CombError;
8use num_bigint::BigUint;
9
10/// Iterator over `k`-combinations of `{0, ..., n-1}` in lexicographic order.
11#[derive(Debug, Clone)]
12pub struct CombinationIter {
13    n: usize,
14    k: usize,
15    current: Option<Vec<usize>>,
16}
17
18/// Construct a combination iterator. `k` must not exceed `n`.
19///
20/// # Examples
21///
22/// The `2`-combinations of `{0, 1, 2}` are produced in lexicographic order:
23///
24/// ```
25/// use sim_lib_discrete_comb::combinations;
26///
27/// let all: Vec<Vec<usize>> = combinations(3, 2).unwrap().collect();
28/// assert_eq!(all, vec![vec![0, 1], vec![0, 2], vec![1, 2]]);
29///
30/// // `k > n` is rejected.
31/// assert!(combinations(2, 3).is_err());
32/// ```
33pub fn combinations(n: usize, k: usize) -> Result<CombinationIter, CombError> {
34    if k > n {
35        return Err(CombError::InvalidParameters(format!(
36            "combinations: k={k} > n={n}"
37        )));
38    }
39    Ok(CombinationIter {
40        n,
41        k,
42        current: Some((0..k).collect()),
43    })
44}
45
46impl Iterator for CombinationIter {
47    type Item = Vec<usize>;
48
49    fn next(&mut self) -> Option<Self::Item> {
50        let cur = self.current.clone()?;
51        // Advance to the next lexicographic combination.
52        let mut next = cur.clone();
53        let mut i = self.k;
54        let advanced = loop {
55            if i == 0 {
56                break false;
57            }
58            i -= 1;
59            if next[i] < self.n - self.k + i {
60                next[i] += 1;
61                for j in (i + 1)..self.k {
62                    next[j] = next[j - 1] + 1;
63                }
64                break true;
65            }
66        };
67        self.current = if advanced { Some(next) } else { None };
68        Some(cur)
69    }
70}
71
72fn validate(combo: &[usize], n: usize) -> Result<(), CombError> {
73    for w in combo.windows(2) {
74        if w[0] >= w[1] {
75            return Err(CombError::InvalidParameters(
76                "combination must be strictly ascending".to_string(),
77            ));
78        }
79    }
80    if let Some(&last) = combo.last()
81        && last >= n
82    {
83        return Err(CombError::OutOfRange {
84            value: last.to_string(),
85            bound: n.to_string(),
86        });
87    }
88    Ok(())
89}
90
91/// The lexicographic rank of `combo` among the `k`-combinations of `n`.
92pub fn combination_rank(combo: &[usize], n: usize) -> Result<BigUint, CombError> {
93    validate(combo, n)?;
94    let k = combo.len();
95    let mut rank = BigUint::from(0u32);
96    let mut prev = 0usize;
97    for (i, &c) in combo.iter().enumerate() {
98        for v in prev..c {
99            rank += binomial((n - 1 - v) as u64, (k - 1 - i) as u64);
100        }
101        prev = c + 1;
102    }
103    Ok(rank)
104}
105
106/// The `k`-combination of `n` at lexicographic ordinal `rank`.
107///
108/// Inverse of [`combination_rank`]: ranking then unranking returns the
109/// original combination.
110///
111/// # Examples
112///
113/// ```
114/// use num_bigint::BigUint;
115/// use sim_lib_discrete_comb::{combination_rank, combination_unrank};
116///
117/// // Among the 2-combinations of {0,1,2}, [0,2] sits at ordinal 1.
118/// assert_eq!(combination_unrank(&BigUint::from(1u32), 3, 2).unwrap(), vec![0, 2]);
119///
120/// // Round-trip rank/unrank.
121/// let combo = vec![1usize, 3];
122/// let r = combination_rank(&combo, 5).unwrap();
123/// assert_eq!(combination_unrank(&r, 5, combo.len()).unwrap(), combo);
124/// ```
125pub fn combination_unrank(rank: &BigUint, n: usize, k: usize) -> Result<Vec<usize>, CombError> {
126    let total = binomial(n as u64, k as u64);
127    if rank >= &total {
128        return Err(CombError::OutOfRange {
129            value: rank.to_string(),
130            bound: total.to_string(),
131        });
132    }
133    let mut remaining = rank.clone();
134    let mut combo = Vec::with_capacity(k);
135    let mut v = 0usize;
136    for i in 0..k {
137        loop {
138            let cnt = binomial((n - 1 - v) as u64, (k - 1 - i) as u64);
139            if remaining < cnt {
140                combo.push(v);
141                v += 1;
142                break;
143            }
144            remaining -= cnt;
145            v += 1;
146        }
147    }
148    Ok(combo)
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn lexicographic_order_and_count() {
157        let all: Vec<_> = combinations(4, 2).unwrap().collect();
158        assert_eq!(
159            all,
160            vec![
161                vec![0, 1],
162                vec![0, 2],
163                vec![0, 3],
164                vec![1, 2],
165                vec![1, 3],
166                vec![2, 3],
167            ]
168        );
169        assert_eq!(combinations(5, 3).unwrap().count(), 10);
170    }
171
172    #[test]
173    fn rank_unrank_round_trip() {
174        for (i, c) in combinations(6, 3).unwrap().enumerate() {
175            let r = combination_rank(&c, 6).unwrap();
176            assert_eq!(r, BigUint::from(i as u32));
177            assert_eq!(combination_unrank(&r, 6, 3).unwrap(), c);
178        }
179    }
180
181    #[test]
182    fn empty_combination_is_singleton() {
183        assert_eq!(combinations(5, 0).unwrap().count(), 1);
184    }
185}
sim_lib_discrete_comb/combination.rs

sim_lib_discrete_comb/
combination.rs
