Python：循环可变数量的输入

Question

Is there a way to efficiently feed in a variable number of inputs into a program like itertools.product() when the number of iterables varies based on user input?

For example, the code below runs, but needs to be explicitly defined for each loop:

steps = np.linspace(0,100,21)

if len(elem_list) == 2:
    for phase in itertools.product(steps,steps):
        if round(round(phase[0],10)+round(phase[1],10),10)==100:
            print(phase)

if len(elem_list) == 3:
    for phase in itertools.product(steps,steps,steps):
        if round(round(phase[0],10)+round(phase[1],10)+round(phase[2],10),10)==100:
            print(phase)
            
if len(elem_list) == 4:
    for phase in itertools.product(steps,steps,steps,steps):
        if round(round(phase[0],10)+round(phase[1],10)+round(phase[2],10)+round(phase[3],10),10)==100:
            print(phase)
            
if len(elem_list) == 5:
    for phase in itertools.product(steps,steps,steps,steps,steps):
        if round(round(phase[0],10)+round(phase[1],10)+round(phase[2],10)+round(phase[3],10)+round(phase[4],10),10)==100:
            print(phase)

where elem_list contains a variable number of elements input by the user.

Is there any way to write this more concisely and so it can be generally applied to any length of elem_list? Thanks!

Answer 1

itertools.product accepts a repeat argument exactly for this. Use it along with sum and map .

I'm not sure why you need round (or round(..., 10) for that matter) since steps includes whole numbers (so summing them is guaranteed to give a whole number) but this is the way to go.

steps = np.linspace(0, 100, 21)

for phase in itertools.product(steps, repeat=len(elem_list)):
    if sum(map(round, phase)) == 100:
        print(phase)

Answer 2

As @DeepSpace mentions, repeat= is the perfect way to handle when all the iterated items are the same. And that sounds like what you wants.

In more complicated cases, you can use a list of ranges. And this list can be arbitrarily generated.

list_of_ranges = [range(5), (1, 3, 7), range(10), (True, False)]
itertools.product(*list_of_ranges)