Rust itertools combinations with variable length counts

Question

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.

Answer 1

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]]

Answer 2

It sounds like what you want to do is partition a series of n items into m sets where each set has a predefined length? You can do this with a recursive approach:

Given a series of lengths lengths and the desired list of things items:

• first check if lengths is empty, if so yield an empty list and stop
• pop the first length from lengths and store it in current_length
• for each combination combination of items with length current_length do:
  1. generate a new list remaining_items with all items from items which are not included in combination
  2. do a recursive call of this function with arguments lengths and remaining_items and for each result rest do:
     • yield rest with combination prepended

This will give you a generator which will produce the desired result without any duplicates.

If you are ok with using rust nightly and the itertools library an implementation is:

#![feature(generators, generator_trait)]

use std::{ops::Generator, iter::FromIterator, pin::Pin};
use std::ops::GeneratorState;
use itertools::Itertools;

fn partition_series(sizes: Vec<usize>, items: Vec<u64>) -> impl Iterator<Item = Vec<Vec<u64>>> {
    GeneratorToIterator(move || {
        if sizes.len() == 0 {
            yield vec![];
            return;
        }
        
        let current_size = sizes[0];
        for combination in items.clone().into_iter().combinations(current_size) {
            let remaining_items: Vec<u64> = items
                .clone()
                .into_iter()
                .filter(|n| !combination.contains(n))
                .collect();

            let inner_generator: Box<dyn Iterator<Item = Vec<Vec<u64>>>> = Box::new(partition_series(sizes[1..].into(), remaining_items));
            for mut rest in inner_generator {
                rest.insert(0, combination.clone());
                yield rest;
            } 
        }
    })
}

struct GeneratorToIterator<G>(G);

impl<G> Iterator for GeneratorToIterator<G>
where
    G: Generator<Return = ()>,
{
    type Item = G::Yield;

    fn next(&mut self) -> Option<Self::Item> {
        let me = unsafe { Pin::new_unchecked(&mut self.0) };
        match me.resume(()) {
            GeneratorState::Yielded(x) => Some(x),
            GeneratorState::Complete(_) => None,
        }
    }
}

Which you can call for a series of number from 0 to n via:

fn main() {
    let sizes = vec![3, 2];
    let total_size = sizes.iter().sum::<usize>() as u64;
    let numbers = Vec::from_iter(0..total_size);
    for combination in partition_series(sizes, numbers) {
        println!("{:?}", combination);
    }
}

Which would produce the following output:

[[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]]

Since the rust implementation might be a bit difficult to understand due to the somewhat unwieldy generator ergonomics, here is a python implementation as wel:

from typing import Iterable, List
from itertools import combinations


def main():
    for combination in generate_next([2, 2, 1]):
        print(combination)

def generate_next(sizes: List[int]) -> Iterable[List[List[int]]]:
    total_size = sum(sizes)
    numbers = list(range(total_size))
    yield from generate_next_internal(sizes, numbers)

def generate_next_internal(sizes: List[int], remaining: List[int]) -> Iterable[List[List[int]]]:
    if len(sizes) == 0:
        yield []
        return

    current_size = sizes.pop(0)

    for combination in combinations(list(remaining), current_size):
        new_remaining = [i for i in remaining if i not in combination]
        for rest in generate_next_internal(list(sizes), new_remaining):
            rest.insert(0, list(combination))
            yield rest



if __name__ == '__main__':
    main()