Generate power sets or subset combinations from (indexed) sequences

counting.py

#!/usr/bin/env python3

def main():
    print("powerset:")
    for subset in powerset([0, 1, 2, 3]):
        print(subset)
    print()
    print("combinations:")
    for subset in combinations([0, 1, 2, 3, 4, 5], 3):
        print(subset)


def powerset(items):
    count = len(items)
    def rec(subset):
        if len(subset) == count:
            yield [items[i] for i in range(count) if subset[i]]
        else:
            yield from rec(subset + [False])
            yield from rec(subset + [True])
    yield from rec([])


def combinations(items, size):
    def rec(subset):
        current_size = len(subset)
        if current_size == size:
            yield [items[i] for i in subset]
        else:
            if current_size == 0:
                start = 0
            else:
                start = subset[-1] + 1
            remaining = size - current_size
            # End index is exclusive
            end = len(items) - remaining + 1
            for i in range(start, end):
                yield from rec(subset + [i])
    yield from rec([])


if __name__ == "__main__":
    main()

counting.rs

fn main() {
    let items = vec![String::from("one"), String::from("two"), String::from("3")];
    println!("powerset:");
    with_powerset(items.as_slice(), &mut |s| println!("{s:?}"));
    println!("combinations (3, 1):");
    with_combinations(items.as_slice(), 1, &mut |s| println!("{s:?}"));
    println!("combinations (3, 2):");
    with_combinations(items.as_slice(), 2, &mut |s| println!("{s:?}"));
    println!("combinations (3, 3):");
    with_combinations(items.as_slice(), 3, &mut |s| println!("{s:?}"));

    let items = vec![
        "foo", "bar", "baz", "qux", "fax", "wax", "win", "bam", "jam",
    ];
    println!("combinations (9, 3):");
    with_combinations(items.as_slice(), 3, &mut |s| println!("{s:?}"));
}

trait LengthIndex: std::ops::Index<usize> {
    fn len(&self) -> usize;
}

impl<T> LengthIndex for [T] {
    fn len(&self) -> usize {
        <[T]>::len(self)
    }
}

fn with_powerset<C, T>(items: &C, f: &mut dyn FnMut(&[&T]))
where
    C: LengthIndex<Output = T> + ?Sized,
{
    fn rec<'a, C, T>(
        items: &'a C,
        f: &mut dyn FnMut(&[&T]),
        next_idx: usize,
        result: &mut Vec<&'a T>,
    ) where
        C: LengthIndex<Output = T> + ?Sized,
    {
        if next_idx == items.len() {
            f(&result);
        } else {
            rec(items, f, next_idx + 1, result);
            result.push(items.index(next_idx));
            rec(items, f, next_idx + 1, result);
            result.pop();
        }
    }
    rec(items, f, 0, &mut Vec::with_capacity(items.len()));
}

fn with_combinations<C, T>(items: &C, size: usize, f: &mut dyn FnMut(&[&T]))
where
    C: LengthIndex<Output = T> + ?Sized,
{
    fn rec<'a, C, T>(
        items: &'a C,
        size: usize,
        f: &mut dyn FnMut(&[&T]),
        next_idx: usize,
        result: &mut Vec<&'a T>,
    ) where
        C: LengthIndex<Output = T> + ?Sized,
    {
        // Precondition: items.len() >= size.
        if result.len() == size {
            f(&result);
        } else {
            let items_needed = size - result.len();
            let items_remaining = items.len() - next_idx;
            if items_needed <= items_remaining - 1 {
                // Optimization: only recurse if we'll have enough after skipping this index.
                rec(items, size, f, next_idx + 1, result);
            }
            if items_needed <= items_remaining {
                // Optimization: only recurse if we'll have enough even after adding this index.
                result.push(items.index(next_idx));
                rec(items, size, f, next_idx + 1, result);
            }
            result.pop();
        }
    }
    if items.len() < size {
        // No point recursing if we can't produce any results.
        return;
    }
    rec(items, size, f, 0, &mut Vec::with_capacity(size));
}