mathrustcombinationspermutationrust-itertools

Rust itertools combinations with variable length counts


I want to calculate a vector's combination.

I am able to do it easily using itertools::Itertools:combinations trait like this:

vec![1, 2, 3].iter().combinations(2).for_each(|x| {
    println!("{:?}", x);
});

But I want to specify the combination lengths as well as counts of these lengths. As an example:

values = [0, 1, 2, 3, 4]

# 1 group with a length of 3 and 1 group with a length of 2
len_counts = { 3: 1, 2: 1 }

combinations = [
    [{0, 1, 2}, {3, 4}]
    [{0, 1, 3}, {2, 4}]
    [{0, 1, 4}, {2, 3}]
    [{0, 2, 3}, {1, 4}]
    [{0, 2, 4}, {1, 3}]
    [{0, 3, 4}, {1, 2}]
    [{1, 2, 3}, {0, 4}]
    [{1, 2, 4}, {0, 3}]
    [{1, 3, 4}, {0, 2}]
    [{2, 3, 4}, {0, 1}]
]

I want it to be lazy-loaded and as clean as possible. I tried to get this output for some time but couldn't succeed. Any help is appreciated.

Edit: The order of combinations and data structures used for representing the variables are not important.


Solution

  • After a bunch of thought, I sadly wasn't able to come up with a clean and easy solution.

    Nonetheless, I came up with a solution :) although it's quite messy, I'm afraid :D

    use std::{collections::HashSet, iter::Peekable};
    
    use itertools::{Combinations, Itertools};
    
    // This struct is so we can `HashSet` by reference address.
    // This prevents that `T` needs to be hashable.
    struct GroupedCombinationsValue<'a, T>(&'a T);
    
    impl<'a, T> GroupedCombinationsValue<'a, T> {
        fn new(val: &'a T) -> Self {
            Self(val)
        }
    }
    
    impl<'a, T> std::hash::Hash for GroupedCombinationsValue<'a, T> {
        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
            std::ptr::hash(self.0, state);
        }
    }
    
    impl<'a, T> PartialEq for GroupedCombinationsValue<'a, T> {
        fn eq(&self, other: &Self) -> bool {
            std::ptr::eq(self.0, other.0)
        }
    }
    
    impl<'a, T> Clone for GroupedCombinationsValue<'a, T> {
        fn clone(&self) -> Self {
            Self(self.0)
        }
    }
    
    impl<'a, T> Eq for GroupedCombinationsValue<'a, T> {}
    
    struct GroupedCombinations<'a, T> {
        values: HashSet<GroupedCombinationsValue<'a, T>>,
        leftover_counts: &'a [usize],
        iter: Peekable<Combinations<std::vec::IntoIter<&'a T>>>,
        child_iter: Option<Box<GroupedCombinations<'a, T>>>,
    }
    
    impl<'a, T> GroupedCombinations<'a, T> {
        fn new(values: Vec<&'a T>, counts: &'a [usize]) -> Self {
            let count;
            let leftover_counts;
    
            if counts.len() == 0 {
                count = 0;
                leftover_counts = counts;
            } else {
                count = counts[0];
                leftover_counts = &counts[1..];
            }
    
            let iter = values.clone().into_iter().combinations(count).peekable();
            let values = values
                .into_iter()
                .map(GroupedCombinationsValue::new)
                .collect::<HashSet<_>>();
    
            Self {
                values,
                leftover_counts,
                iter,
                child_iter: None,
            }
        }
    }
    
    impl<'a, T> Iterator for GroupedCombinations<'a, T> {
        type Item = Vec<Vec<&'a T>>;
    
        fn next(&mut self) -> Option<Self::Item> {
            let local_value = self.iter.peek()?;
    
            if self.child_iter.is_none() && !self.leftover_counts.is_empty() {
                let child_values = self
                    .values
                    .difference(
                        &local_value
                            .iter()
                            .cloned()
                            .map(GroupedCombinationsValue::new)
                            .collect(),
                    )
                    .map(|e| e.0)
                    .collect::<Vec<_>>();
                self.child_iter = Some(Box::new(Self::new(child_values, self.leftover_counts)));
            }
    
            let mut result = vec![];
            if !local_value.is_empty() {
                result.extend_from_slice(&[local_value.clone()]);
            }
    
            if let Some(child_iter) = &mut self.child_iter {
                match child_iter.next() {
                    Some(child_value) => {
                        result.extend(child_value);
                        Some(result)
                    }
                    None => {
                        self.child_iter = None;
                        self.iter.next();
                        self.next()
                    }
                }
            } else {
                self.iter.next();
                Some(result)
            }
        }
    }
    
    fn grouped_combinations<'a, T>(values: &'a [T], counts: &'a [usize]) -> GroupedCombinations<'a, T> {
        GroupedCombinations::new(values.iter().collect(), counts)
    }
    
    fn main() {
        let values = [0, 1, 2, 3, 4];
        let counts = [3, 2];
    
        let combinations = grouped_combinations(&values, &counts);
    
        for combination in combinations {
            println!("{:?}", combination);
        }
    }
    
    [[0, 1, 2], [3, 4]]
    [[0, 1, 3], [2, 4]]
    [[0, 1, 4], [2, 3]]
    [[0, 2, 3], [1, 4]]
    [[0, 2, 4], [1, 3]]
    [[0, 3, 4], [1, 2]]
    [[1, 2, 3], [4, 0]]
    [[1, 2, 4], [3, 0]]
    [[1, 3, 4], [2, 0]]
    [[2, 3, 4], [1, 0]]