Finding all non-negative integers that add up to a particular value or less, with number of integers-CodePudding

I am working in Python. I want to make a function that takes two values: maxSum and n. The output should be all the n-tuples of non-negative integers that add up to maxSum or less. For example, if maxSum is 2 and n is three, the output should be:

[[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [0, 1, 1], [0, 2, 0], [1, 0, 0], [1, 0, 1], [1, 1, 0], [2, 0, 0]]

For any given value of n, I can do it, just by going through n for-loops. For example, if n is 3, I can simply do:

result = []

for i in range(maxSum   1):
    for j in range(maxSum   1 - i):
        for k in range(maxSum   1 - i - j):
            tuple = [i, j, k]
            result.insert(len(result), tuple)
return result

But in the above example, I have typed the three for-loops by hand. That is not what I want. I want to be able to make a function that works for any n, that is, a function tupleGenerator(maxSum, n) that generates all such n-tuples. The problem is, I am unable to change the number of for-loops in the function based on the input n.

CodePudding user response：

Sorry don't have enough rep yet to comment, but here's a solution from math stackexchange

https://math.stackexchange.com/questions/2410399/how-to-find-all-the-n-tuples-with-nonnegative-integers-that-add-up-to-a-fixed-in

In particular Ocean Wong's answer, which includes a recursive Python solution (last answer on the page).

EDIT: slightly revised answer to include code as OP requested.

def stars_and_bars(num_objects_remaining, num_bins_remaining, filled_bins=()):
    """
    The Stars and Bars (https://en.wikipedia.org/wiki/Stars_and_bars_(combinatorics)) problem can be thought of as
    1. Putting m-1 bars ("|") in amongst a row of n stars
    OR
    2. Putting n balls into m bins.

    This program lists all the possible combinations by doing so.
    """
    if num_bins_remaining > 1:
        # case 1: more than one bins left
        novel_arrangement = []  # receptacle to contain newly generated arrangements.
        for num_distributed_to_the_next_bin in range(0, num_objects_remaining   1):
            # try putting in anything between 0 to num_objects_remaining of objects into the next bin.
            novel_arrangement.extend(
                stars_and_bars(
                    num_objects_remaining - num_distributed_to_the_next_bin,
                    num_bins_remaining - 1,
                    filled_bins=filled_bins   (num_distributed_to_the_next_bin,),
                )
                # one or multiple tuple enclosed in a list
            )
        return novel_arrangement
    else:
        # case 2: reached last bin. terminate recursion.
        # return a single tuple enclosed in a list.
        return [
            filled_bins   (num_objects_remaining,),
        ]

result = []
n = 3
for i in range(maxSum   1):
    result.extend(stars_and_bars(i, n))
print(result)

CodePudding user response：

Python provides the itertools module that you can leverage for this:

maxSum, n = 2, 3
p1 = it.combinations_with_replacement(range(n), n)
p2 = filter( lambda m: sum(m) <= maxSum, p1 )
p3 = it.chain.from_iterable(it.permutations(p) for p in p2)
p4 = list(set(list(p3)))
sorted(p4)

This should be much easier on the memory consumption. You should print each step (for example print(list(p1))) to understand what each step is doing. That is easier to understand than anything that I will write in text here.