Iterating over every vector of N non-negative integers which adds up to X

Hello!
I'm fairly new to rust and I'm practicing it by solving Advent of Code 2015. So far I've arrived at day 15, an optimization problem where the problem space is a set of N parameters (the ingredients) that have to add up to a constant (100, in this case). As a first step I have to write a function that takes those two parameters and results in a iterator of Vecs, each of length N

For example, if N is 2 the combinations are [100,0], [99,1], ..., [0,100]. The problem scales rapidly with N, as each new parameter adds a dimension to the space.

I can't find a concise way to phrase this in a search engine, so I might have missed an already existing answer, but before i try to implement the solution completely manually, does the rust standard library or an external crate contain something that might help accomplish this task?

Thank you for the attention

After looking at the problem description, I’ve noticed that your title inaccurately described the problem statement. To avoid similar confusion for future readers, I’ve adapted it for you.

In particular: The ordering matters here, in particular considering [0, 100] separately from [100, 0] is important, so it’s not about finding all “sets” but all “lists” or “tuples” or “vectors”, whichever term you prefer. As your post contents mentions “Vecs”, I’ve gone with “vectors” for the title.

Second, “positive integer” usually refers to integers strictly greater than 0, i.e. not including zero, whereas as far as I understand in this AoC problem, amount zero is allowed for an ingredient.

2 Likes

Regarding solution, there’s really two straightforward approaches: The first one is to iterate over all the length N vecs with entries between 0 and 100, and filter out those that don’t add up. The second one is to recursively generate these while paying attention to the remaining capacity in order to ensure the summation to 100 by construction. The latter approach is more efficient, however it’s not terribly more efficient, as both approaches will be processing more-or-less exponentially many elements, measured in terms of N.

The first approach is somewhat more straightforward to implement. You can use existing infrastructure from itertools, in particular Itertools::multi_cartesian_product:

Code for one solution using multi_cartesian_product and filter (click to expand)
use itertools::Itertools;

fn all_combinations(n: usize) -> impl Iterator<Item = Vec<u32>> {
    itertools::repeat_n(0..=100, n)
        .multi_cartesian_product()
        .filter(|v| v.iter().sum::<u32>() == 100)
}

fn main() {
    for v in all_combinations(3) {
        println!("{v:?}");
    }
}
2 Likes

One possible way this could be done is by recursively defining a function

fn all_combinations(n: u32, cap: u32) -> Box<dyn Iterator<Item = Vec<u32>>>

(the main call would then start by passing cap: 100; e.g. all_combinations(3, 100) to generate the same output as the code in my previous answer)

This function can be defined recursively fairly straightforward using flat_map and map from the Iterator trait.

Further hints (feel free not to read these if you want to try this yourself):

  • The base case is n == 0, in which case you can return Box::new(iter::once(vec![cap])).
  • The (non base case) all_combinations(n, cap) implementation works by recursively calling all_combinations(n - 1, …) multiple times (with different values for the “”)
  • Use move closures
  • A helper function to add an element to a Vec could be helpful for cleaner code, e.g. a function like
    fn cons<T: Copy>(x: T, xs: &[T]) -> Vec<T>
    
Full code (click to expand)
use std::iter;

fn cons<T: Copy>(x: T, xs: &[T]) -> Vec<T> {
    let mut v = Vec::with_capacity(xs.len() + 1);
    v.push(x);
    v.extend(xs);
    v
}

fn all_combinations(n: u32, cap: u32) -> Box<dyn Iterator<Item = Vec<u32>>> {
    assert_ne!(n, 0);
    if n == 1 {
        Box::new(iter::once(vec![cap]))
    } else {
        Box::new((0..=cap).flat_map(move |x| all_combinations(n - 1, cap - x).map(move |xs| cons(x, &xs))))
    }
}

fn main() {
    for v in all_combinations(3, 100) {
        println!("{v:?}");
    }
}

P.S. these are the kinds of algorithms where I still prefer writing Haskell :D
import Data.Foldable

allCombinations n _ | n < 1 = error "n must be at least 1"
allCombinations 1 cap = [[cap]]
allCombinations n cap = [ x:xs | x <- [0..cap], xs <- allCombinations (n-1) (cap - x)]

main = for_ (allCombinations 3 100) $ \v ->
    print v

(compiler explorer)

1 Like

I'm sorry I can't mark both of your code replies as a solution. If this was production code I'd go with the recursive approach, but since it has to run only once the "brute force" one is easier to implement. I'll post an update as soon as i get it working, thanks for the assistance!

For production code, you can go further and avoid a lot of allocations by using a callback instead of producing an iterator.

fn for_all_combinations_rec(n: usize, cap: u32, vec: &mut Vec<u32>, f: &mut impl FnMut(&[u32])) {
    assert_ne!(n, 0);
    if n == 1 {
        vec.push(cap);
        f(vec);
        vec.pop();
    } else {
        (0..=cap).for_each(|x| {
            vec.push(x);
            for_all_combinations_rec(n - 1, cap - x, vec, f);
            vec.pop();
        })
    }
}

fn for_all_combinations(n: usize, cap: u32, mut f: impl FnMut(&[u32])) {
    for_all_combinations_rec(n, cap, &mut Vec::with_capacity(n), &mut f)
}

fn main() {
    for_all_combinations(3, 100, |v| {
        println!("{v:?}");
    });
}
2 Likes

In terms of the twentyfold way, this is a multicombination (or n-multisubset): You have 100 unlabeled balls (ingredients) that you want to distribute into N labeled buckets (recipes) with no restriction on the counts. For 6 recipes, there about 108 multicombinations, but there are over 1012 ways to "brute force" the problem, so I suggest the recursive approach even for online challenges like these. Possibly they won't give you so many recipes it's a concern, though.

Or depending on the challenge, they may give you so many that you need to thin the search space in other ways even when you don't generate invalid states, e.g. having to discard large swaths of possible vectors due to the constraints around negative numbers so that you avoid billions of visits you would otherwise check.

I feel the recursive approach is rather understandable, and it's also pretty amenable to culling the problem space when necessary, since you can just not let a given index exceed a problematic cap, say. If I were doing the challenge, I'd probably just use that.


That's all the practical input I have.

However, just for fun, this seemed like a good chance to try implementing a more optimized algorithm in Rust. I remembered reading about a pseduo-Gray code for n-choose-k combinations, where a bit string representing "chosen or not" changes minimally between each visit. (As discussed in the multicombination link, a multicombination can be modeled as a n-choose-k combination.) "Pseudo" as at least two bits much change per visit in order to maintain the set bit count. It seemed like a nice fit since Rust supports 128-bit integers, and 28 recipes would result in too many multicombinations to visit them all, so the limitations are reasonable.

I chose Algorithm C for Chase's sequence from TAOCP [1]. It uses bookkeeping instead of recursion; the generated sequence has enough structure that some sort of culling would be possible, but I didn't explore that at all. The algorithm is pretty much just directly translated from the book. Upon each visit, 1 ingredient moves from one recipe to another (one count goes down by 1 and another goes up by 1).

Again, I still recommend the recursive approach; it's pretty intuitive to wrap your head around and to understand the sequence order too, and how to jump around in it if needed. But I had fun and thought I'd share.


  1. The Art of Computer Programming Volume 4A Part 1, section 7.2.1.3 ↩︎

2 Likes

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.