Search code examples
pythonloopsnested-loops

Number of nested loops by variable

I want my program to find the combinations of x integers with values in a range that summed equal a given number.

For example: I want to find the combinations of 3 integers that can have a value between 0 and 2 and that summed equal 5. I can do this by coding:

possibilities = []
total = 5
valueRange = 3
for num1 in xrange(valueRange):
    for num2 in xrange(valueRange):
        for num3 in xrange(valueRange):
            if num1 + num2 + num3 == total:
                possibilities.append([num1, num2, num3])

I can change the value of the total sum and the range value by changing the variables I created, but how can I specify the amount of nested loops using a variable? Thank you

Solution

  • The itertools.product() function should help:

    >>> [values for values in product(range(3), repeat=3) if sum(values) == 5]
[(1, 2, 2), (2, 1, 2), (2, 2, 1)]

    Seeing that the answer contains anagrams of the same answer, you can reduce the work further by using itertools.combinations_with_replacement():

    >>> # 4 digits in [0, 1, 2, 3, 4] summing to 6
>>> for values in combinations_with_replacement(range(5), 4):
        if sum(values) == 6:
            print(values)

(0, 0, 2, 4)
(0, 0, 3, 3)
(0, 1, 1, 4)
(0, 1, 2, 3)
(0, 2, 2, 2)
(1, 1, 1, 3)
(1, 1, 2, 2)