Python: find all combinations of elements in list to reach a certain value

Question

I'm given a list, for example: [2,5,8] and a value c , for example, 33. So, because the length of the list is 3 , I have to return a list of tuples (i,j,k) for which i*w[0] + j*w[1] + k*w[2] == c (33) .

Of course, this is easily solved with 3 for loops, as shown:

res = []
for i in range(c+1):
   for j in range(c+1):
       for g in range(c+1):
           if i*self.w[0] + j*self.w[1] + g*self.w[2] == c:
              res.append((i,j,g))
return sorted(res)

The output will be:

[(0, 5, 1), (1, 3, 2), (2, 1, 3), (4, 5, 0), (5, 3, 1), (6, 1, 2), (9, 3, 0), (10, 1, 1), (14, 1, 0)]

The problem is that I have to do this with lists with other lengths as well. So my question is; does anyone know how to do this in general, for a list with length x (not too big) and a given value c ?

Answer 1

Using itertools.product and operator.mul :

>>> [k for k in product(range(c+1), repeat=len(w)) if sum(map(mul, w, k)) == c]
[(0, 5, 1), (1, 3, 2), (2, 1, 3), (4, 5, 0), (5, 3, 1), (6, 1, 2), (9, 3, 0), (10, 1, 1), (14, 1, 0)]

Answer 2

You can use itertools.product and zip to generalise here:

import itertools
w = [3,1,2,5]
c = 33
l = [range(c + 1)] * len(w)
res = []
for values in itertools.product(*l):
    if sum(a * b for a, b in zip(w, values)) == c:
        res.append(values)
print(sorted(res))

The loop that builds up res could be replaced by a list comprehension if desired. That is to say:

res = [values for values in itertools.product(*l)
           if sum(a * b for a, b in zip(w, values)) == c]