Skip to main content

sim_lib_discrete_comb/
subset.rs

1//! Subsets of `{0, ..., n-1}` in bitmask order, with rank/unrank.
2//!
3//! A subset is represented as its sorted member indices. Ordinal `r` is the
4//! bitmask whose bit `i` indicates membership of element `i`.
5
6use crate::error::CombError;
7use num_bigint::BigUint;
8
9const MAX_N: usize = 127;
10
11/// Iterator over all `2^n` subsets of `{0, ..., n-1}` in bitmask order.
12#[derive(Debug, Clone)]
13pub struct SubsetIter {
14    n: usize,
15    next: u128,
16    total: u128,
17}
18
19/// Construct a subset iterator, rejecting `n` too large for the `u128` cursor.
20pub fn subsets(n: usize) -> Result<SubsetIter, CombError> {
21    if n > MAX_N {
22        return Err(CombError::LimitExceeded(format!(
23            "subset cardinality {n} exceeds {MAX_N}"
24        )));
25    }
26    let total = if n == MAX_N { u128::MAX } else { 1u128 << n };
27    Ok(SubsetIter { n, next: 0, total })
28}
29
30impl Iterator for SubsetIter {
31    type Item = Vec<usize>;
32
33    fn next(&mut self) -> Option<Self::Item> {
34        if self.next >= self.total {
35            return None;
36        }
37        let r = self.next;
38        self.next += 1;
39        Some((0..self.n).filter(|&i| (r >> i) & 1 == 1).collect())
40    }
41}
42
43/// The bitmask ordinal of `subset` (a list of distinct indices `< n`).
44pub fn subset_rank(subset: &[usize], n: usize) -> Result<BigUint, CombError> {
45    let mut rank = BigUint::from(0u32);
46    for &i in subset {
47        if i >= n {
48            return Err(CombError::OutOfRange {
49                value: i.to_string(),
50                bound: n.to_string(),
51            });
52        }
53        rank.set_bit(i as u64, true);
54    }
55    Ok(rank)
56}
57
58/// The subset (sorted member indices) for bitmask ordinal `rank` over `n`.
59pub fn subset_unrank(rank: &BigUint, n: usize) -> Vec<usize> {
60    (0..n).filter(|&i| rank.bit(i as u64)).collect()
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn powerset_count_and_order() {
69        let all: Vec<_> = subsets(3).unwrap().collect();
70        assert_eq!(all.len(), 8);
71        assert_eq!(all[0], Vec::<usize>::new());
72        assert_eq!(all[1], vec![0]);
73        assert_eq!(all[7], vec![0, 1, 2]);
74    }
75
76    #[test]
77    fn rank_unrank_round_trip() {
78        for (i, s) in subsets(4).unwrap().enumerate() {
79            let r = subset_rank(&s, 4).unwrap();
80            assert_eq!(r, BigUint::from(i as u32));
81            assert_eq!(subset_unrank(&r, 4), s);
82        }
83    }
84
85    #[test]
86    fn rank_rejects_out_of_range() {
87        assert!(matches!(
88            subset_rank(&[5], 4),
89            Err(CombError::OutOfRange { .. })
90        ));
91    }
92}