misc_iterators 0.3.3

/// Iterate over all tuples $(x_1, x_2, ..., x_k)$ of positive integers,
/// such that $\sum^k_{i=1} x_i = n$ for some positive integer $n$.
#[derive(Clone)]
pub struct ExactTupleSum<const K: usize>
{
    // TODO(low): replace with generic const expression when stable.
    indices: [usize; K],
    upper_bound: usize,
}

impl<const K: usize> ExactTupleSum<K>
{
    /// Construct an ExactTupleSum iterator with $K$ partitions which
    /// will add to exactly `n`.
    pub fn new(n: usize) -> Self
    {
        Self {
            indices: [0; K],
            upper_bound: n,
        }
    }
}

impl<const K: usize> Iterator for ExactTupleSum<K>
{
    type Item = [usize; K];

    fn next(&mut self) -> Option<Self::Item>
    {
        // TODO(medium): imporve this by storing the limits and
        // the current index to update, instead of recomputing them each time.
        let mut limits = vec![self.upper_bound, K - 1];
        for i in 1..K - 1
        {
            limits[i] = limits[i - 1] - self.indices[i - 1];
        }

        for i in (0..K - 1).rev()
        {
            if self.indices[i] < limits[i]
            {
                self.indices[i] += 1;
                break;
            }
            else if i > 0
            {
                self.indices[i] = 0;
            }
            else
            {
                return None;
            }
        }

        if self.indices[0] <= self.upper_bound
        {
            let mut res = self.indices;
            res[K - 1] = self.upper_bound;
            for i in 0..K - 1
            {
                res[K - 1] -= res[i];
            }
            return Some(res);
        }

        None
    }
}

/// Iterate over all tuples $(x_1, x_2, ..., x_k)$ of positive integers,
/// such that $\sum^k_{i=1} x_i \leq n$ for some positive integer $n$.
#[derive(Clone)]
pub struct BoundedTupleSum<const K: usize>
{
    // TODO(low): replace with generic const expression when stable.
    bar_positions: Vec<usize>,
    upper_bound: usize,
}

impl<const K: usize> BoundedTupleSum<K>
{
    /// Construct a BoundedTupleSum iterator with $K$ partitions which
    /// will add to at most `n`.
    pub fn new(n: usize) -> Self
    {
        let mut positions: Vec<_> = (0..K + 1).collect();
        positions[K] = n + K;
        Self {
            bar_positions: positions,
            upper_bound: n,
        }
    }

    /// Modify the starting positions, effectively skipping all prior iterations
    /// that would lead to this configuration. `positions` must be a strictly
    /// increasing array of integers.
    pub fn start_at(mut self, positions: [usize; K]) -> Self
    {
        debug_assert!(positions.windows(2).all(|w| w[0] < w[1]));

        for (i, v) in positions.into_iter().enumerate()
        {
            self.bar_positions[i] = v;
        }

        self
    }

    /// Given a tuple whose components add to less than an implcit `n`
    /// generate the index this iterator would have generated for that tuple.
    pub fn to_index(tuple: &[usize; K]) -> usize
    {
        let mut offset = tuple[0];
        let mut cumulative_sum = offset;
        for i in 1..K
        {
            cumulative_sum += 1 + tuple[i];
            offset += num_integer::binomial(cumulative_sum, i + 1);
        }

        offset
    }

    /// Total number of entries generated for a decomposition of size `size`.
    pub fn total(size: usize) -> usize { num_integer::binomial(size + K, K) }
}

impl<const K: usize> ExactSizeIterator for BoundedTupleSum<K>
{
    fn len(&self) -> usize { Self::total(self.upper_bound) }
}

impl<const K: usize> Iterator for BoundedTupleSum<K>
{
    type Item = [usize; K];

    fn next(&mut self) -> Option<Self::Item>
    {
        let construct_tuple = |array: &[usize]| {
            let mut result = [0; K];
            let mut lower_bound = 0;
            let mut correction = 0;
            for i in 0..array.len()
            {
                result[i] = array[i] - lower_bound - correction;
                correction = 1;
                lower_bound = array[i];
            }

            result
        };

        let position =
            (0..K).position(|i| self.bar_positions[i] + 1 < self.bar_positions[i + 1]);

        if let Some(index) = position
        {
            let tuple = construct_tuple(&self.bar_positions[0..K]);

            for i in 0..index
            {
                self.bar_positions[i] = i;
            }

            self.bar_positions[index] += 1;

            return Some(tuple);
        }

        None
    }
}

#[cfg(test)]
mod tests
{
    use super::*;

    #[test]
    fn test_integer_decomposition_to_index()
    {
        let index = BoundedTupleSum::<3>::to_index(&[3, 0, 0]);
        assert_eq!(index, 19);

        let index = BoundedTupleSum::<3>::to_index(&[2, 1, 0]);
        assert_eq!(index, 18);

        let index = BoundedTupleSum::<3>::to_index(&[0, 0, 1]);
        assert_eq!(index, 1);

        let index = BoundedTupleSum::<3>::to_index(&[0, 2, 0]);
        assert_eq!(index, 7);
    }

    #[test]
    fn test_bounded_sum()
    {
        let iter = BoundedTupleSum::<3>::new(3);
        let table = [
            [0, 0, 0],
            [0, 0, 1],
            [0, 1, 0],
            [1, 0, 0],
            [0, 0, 2],
            [0, 1, 1],
            [1, 0, 1],
            [0, 2, 0],
            [1, 1, 0],
            [2, 0, 0],
            [0, 0, 3],
            [0, 1, 2],
            [1, 0, 2],
            [0, 2, 1],
            [1, 1, 1],
            [2, 0, 1],
            [0, 3, 0],
            [1, 2, 0],
            [2, 1, 0],
            [3, 0, 0],
        ];
        for (i, tuple) in iter.enumerate()
        {
            assert_eq!(tuple, table[i]);
        }
    }

    #[test]
    fn test_exact_sum()
    {
        let iter = ExactTupleSum::<3>::new(3);
        let table = [
            [0, 1, 2],
            [0, 2, 1],
            [0, 3, 0],
            [1, 0, 2],
            [1, 1, 1],
            [1, 2, 0],
            [2, 0, 1],
            [2, 1, 0],
            [3, 0, 0],
        ];
        for (i, tuple) in iter.enumerate()
        {
            assert_eq!(tuple, table[i]);
        }
    }
}