1
\$\begingroup\$

I'm doing some boring math with lot of sums over lot of variables. So, instead of nesting for loops I have decided to use something like this:

def combo(*limits):
    if len(limits) == 1:
        for l in range(limits[0]):
            yield [l]
    else:
        for l in range(limits[0]):
            for l2 in combo(*limits[1:]):
                yield [l] + l2


if __name__=="__main__":
    for i, j, k in combo(2, 4, 5):
        print i, j, k

Do you like it? Could it be redone or should it be replaced by something else?

\$\endgroup\$
5
\$\begingroup\$

Many interesting loops in Python are better written using itertools. This problem is basically just a kind of product().

from itertools import product

def combo(*limits):
    return product(*map(range, limits))
\$\endgroup\$
3
\$\begingroup\$

I suggest to rethink your parameter list. As it stand, if someone tries to call this function without parameters (which is perfectly valid looking at the signature), your function crashes with an IndexError which is not very friendly. Instead, you should declare def combo(term, *remainders): to make it obvious that the function need at least 1 argument. And when call it with 0 arguments, the error will be more explicit:

TypeError: combo() missing 1 required positional argument: 'term'

It also lets you test for the end of the recursion more easily:

def combo(term, *remainders):
    if not remainders:
        for head in range(term):
            yield [head]
    else:
        for head in range(term):
            for tail in combo(*remainders):
                yield [head] + tail

You could also put the test under the first for loop to avoid repeating it twice…

\$\endgroup\$
  • \$\begingroup\$ I like that one! \$\endgroup\$ – Eugene Feb 11 '16 at 9:37
  • \$\begingroup\$ sorry for dis-accepting your answer, but 200_success's answer is better. Still thanks for the parameter hint \$\endgroup\$ – Eugene Feb 11 '16 at 10:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.