Any way to speedup itertool.product

Question

I am using itertools.product to find the possible weights an asset can take given that the sum of all weights adds up to 100.

min_wt = 10
max_wt = 50
step = 10
nb_Assets = 5

weight_mat = []
for i in itertools.product(range(min_wt, (max_wt+1), step), repeat = nb_Assets):
    if sum(i) == 100:
        weight = [i]
        if np.shape(weight_mat)[0] == 0:
            weight_mat = weight
        else:
            weight_mat = np.concatenate((weight_mat, weight), axis = 0)

The above code works, but it is too slow as it goes through the combinations that are not acceptable, example [50,50,50,50,50] eventually testing 3125 combinations instead of 121 possible combinations. Is there any way we can add the 'sum' condition within the loop to speed things up?

Answer 1

Many improvements are possible.

For starters, the search space can be reduced using itertools.combinations_with_replacement() because summation is commutative.

Also, the last addend should be computed rather than tested. For example if t[:4] was (10, 20, 30, 35), you could compute t[4] as 1 - sum(t), giving a value of 5. This will give a 100-fold speed-up over trying one-hundred values of x in (10, 20, 30, 35, x).

Answer 2

You can write up a recursive algorithm for that which prunes all the impossible options early on:

def make_weight_combs(min_wt, max_wt, step, nb_assets, req_wt):
    weights = range(min_wt, max_wt + 1, step)
    current = []
    yield from _make_weight_combs_rec(weights, nb_assets, req_wt, current)

def _make_weight_combs_rec(weights, nb_assets, req_wt, current):
    if nb_assets <= 0:
        yield tuple(current)
    else:
        # Discard weights that cannot possibly be used
        while weights and weights[0] + weights[-1] * (nb_assets - 1) < req_wt:
            weights = weights[1:]
        while weights and weights[-1] + weights[0] * (nb_assets - 1) > req_wt:
            weights = weights[:-1]
        # Add all possible weights
        for w in weights:
            current.append(w)
            yield from _make_weight_combs_rec(weights, nb_assets - 1, req_wt - w, current)
            current.pop()

min_wt = 10
max_wt = 50
step = 10
nb_assets = 5
req_wt = 100
for comb in make_weight_combs(min_wt, max_wt, step, nb_assets, req_wt):
    print(comb, sum(comb))

Output:

(10, 10, 10, 20, 50) 100
(10, 10, 10, 30, 40) 100
(10, 10, 10, 40, 30) 100
(10, 10, 10, 50, 20) 100
(10, 10, 20, 10, 50) 100
(10, 10, 20, 20, 40) 100
(10, 10, 20, 30, 30) 100
(10, 10, 20, 40, 20) 100
...

If order of the weights does not matter (so, for example, (10, 10, 10, 20, 50) and (50, 20, 10, 10, 10) are the same), then you can modify the for loop as follows:

for i, w in enumerate(weights):
    current.append(w)
    yield from _make_weight_combs_rec(weights[i:], nb_assets - 1, req_wt - w, current)
    current.pop()

Which gives the output:

(10, 10, 10, 20, 50) 100
(10, 10, 10, 30, 40) 100
(10, 10, 20, 20, 40) 100
(10, 10, 20, 30, 30) 100
(10, 20, 20, 20, 30) 100
(20, 20, 20, 20, 20) 100

Answer 3

Let's generalise this problem; you want to iterate over k-tuples whose sum is n, and whose elements are within range(min_w, max_w+1, w_step). This is a kind of integer partitioning problem, with some extra constraints on the size of the partition and the sizes of its components.

To do this, we can write a recursive generator function; for each w in the range, the remainder of the tuple is a (k - 1)-tuple whose sum is (n - w). The base case is a 0-tuple, which is possible only if the required sum is 0.

As Raymond Hettinger notes, you can also improve the efficiency when k = 1 by just testing whether the required sum is one of the allowed weights.

def constrained_partitions(n, k, min_w, max_w, w_step=1):
    if k < 0:
        raise ValueError('Number of parts must be at least 0')
    elif k == 0:
        if n == 0:
            yield ()
    elif k == 1:
        if n in range(min_w, max_w+1, w_step):
            yield (n,)
    elif min_w*k <= n <= max_w*k:
        for w in range(min_w, max_w+1, w_step):
            for p in constrained_partitions(n-w, k-1, min_w, max_w, w_step):
                yield (w,) + p

Usage:

>>> for p in constrained_partitions(5, 3, 1, 5, 1):
...     print(p)
...
(1, 1, 3)
(1, 2, 2)
(1, 3, 1)
(2, 1, 2)
(2, 2, 1)
(3, 1, 1)
>>> len(list(constrained_partitions(100, 5, 10, 50, 10)))
121

Whenever you're iterating over all solutions to some sort of combinatorial problem, it's generally best to generate actual solutions directly, rather than generate more than you need (e.g. with product or combinations_with_replacement) and reject the ones you don't want. For larger inputs, the vast majority of time would be spent generating solutions which will get rejected, due to combinatorial explosion.

Note that if you don't want repeats in different orders (e.g. 1, 1, 3 and 1, 3, 1), you can change the recursive call to constrained_partitions(n-w, k-1, min_w, w, w_step) to only generate partitions where the weights are in non-increasing order.

Answer 4

Comparing performance of the offered solutions:

import itertools
import timeit
import numpy as np


# original code from question
def f1():
    min_wt = 10
    max_wt = 50
    step = 10
    nb_assets = 5

    weight_mat = []
    for i in itertools.product(range(min_wt, (max_wt+1), step), repeat=nb_assets):
        if sum(i) == 100:
            weight = [i, ]
            if np.shape(weight_mat)[0] == 0:
                weight_mat = weight
            else:
                weight_mat = np.concatenate((weight_mat, weight), axis=0)

    return weight_mat


# code from question using list instead of numpy array
def f1b():
    min_wt = 10
    max_wt = 50
    step = 10
    nb_assets = 5

    weight_list = []
    for i in itertools.product(range(min_wt, (max_wt+1), step), repeat=nb_assets):
        if sum(i) == 100:
            weight_list.append(i)

    return weight_list


# calculating the last element of each tuple
def f2():
    min_wt = 10
    max_wt = 50
    step = 10
    nb_assets = 5

    weight_list = []
    for i in itertools.product(range(min_wt, (max_wt+1), step), repeat=nb_assets-1):
        the_sum = sum(i)
        if the_sum < 100:
            last_elem = 100 - the_sum
            if min_wt <= last_elem <= max_wt:
                weight_list.append(i + (last_elem, ))

    return weight_list


# recursive solution from user kaya3 (https://stackoverflow.com/a/58823843/9225671)
def constrained_partitions(n, k, min_w, max_w, w_step=1):
    if k < 0:
        raise ValueError('Number of parts must be at least 0')
    elif k == 0:
        if n == 0:
            yield ()
    else:
        for w in range(min_w, max_w+1, w_step):
            for p in constrained_partitions(n-w, k-1, min_w, max_w, w_step):
                yield (w,) + p

def f3():
    return list(constrained_partitions(100, 5, 10, 50, 10))


# recursive solution from user jdehesa (https://stackoverflow.com/a/58823990/9225671)
def make_weight_combs(min_wt, max_wt, step, nb_assets, req_wt):
    weights = range(min_wt, max_wt + 1, step)
    current = []
    yield from _make_weight_combs_rec(weights, nb_assets, req_wt, current)

def _make_weight_combs_rec(weights, nb_assets, req_wt, current):
    if nb_assets <= 0:
        yield tuple(current)
    else:
        # Discard weights that cannot possibly be used
        while weights and weights[0] + weights[-1] * (nb_assets - 1) < req_wt:
            weights = weights[1:]
        while weights and weights[-1] + weights[0] * (nb_assets - 1) > req_wt:
            weights = weights[:-1]
        # Add all possible weights
        for w in weights:
            current.append(w)
            yield from _make_weight_combs_rec(weights, nb_assets - 1, req_wt - w, current)
            current.pop()

def f4():
    return list(make_weight_combs(10, 50, 10, 5, 100))

I tested these functions using timeit like this:

print(timeit.timeit('f()', 'from __main__ import f1 as f', number=100))

The results using the parameters from the question:

# min_wt = 10
# max_wt = 50
# step = 10
# nb_assets = 5
0.07021828400320373       # f1 - original code from question
0.041302188008558005      # f1b - code from question using list instead of numpy array
0.009902548001264222      # f2 - calculating the last element of each tuple
0.10601829699589871       # f3 - recursive solution from user kaya3
0.03329997700348031       # f4 - recursive solution from user jdehesa

If I expand the search space (reduced step and increased assets):

# min_wt = 10
# max_wt = 50
# step = 5
# nb_assets = 6
7.6620834979985375        # f1 - original code from question
7.31425816299452          # f1b - code from question using list instead of numpy array
0.809070186005556         # f2 - calculating the last element of each tuple
14.88188026699936         # f3 - recursive solution from user kaya3
0.39385621099791024       # f4 - recursive solution from user jdehesa

Seems like f2 and f4 are the fastest (for the tested size of the data).

Answer 5

Note that when you have N weights that sum up to 100, and you chose N - 1 weights, the remaining weight is already defined as 100 - sum of already chosen weights, which should be positive. The same limitation applies to any number of already chosen weights.

Next, you don't want combinations that are just permutations of the same weights. This is why you can order weights by value, and choose the next weight in the combination to be below or equal of the previous one.

This immediately makes the search space much smaller, and you can break a particular branch of search earlier.

Probably writing it with explicit loops first, or as a recursive algorithm, should be much easier for understanding and implementing.