Search code examples
pythonfor-loopnested-loops

Reducing nested for loops that build dictionaries/strings (Branching tree style)


For a recent script I had to iterate possible flows through a branching tree. I built a dictionary that would contain the run state for each item in order for me to evaluate. For example:

for a in range(0, 2):
  for b in range(0, 2):
    for c in range(0, 2):
      for d in range(0, 2):
        ...
        run = [a,b,c,d ...]

Unfortunately, it started out small but grew to over a dozen states. How could this be reduced to eliminate all of the nested loops?

Would a similar answer differ if some of the states had 3 or 4 states instead of 2?

Similarly, how would the same issue be affected if each loop was from a function list? I suspected it would be the same. For example:

def leet_lookup(char):
    table = {"a": ["a","A","@"],
            "b": ["b", "B", "8"],
            "c": ["c", "C", "<"],
            "e": ["e", "E", "3"],
            "i": ["i", "I", "1"],
            "l": ["l", "L", "1"],
            "o": ["o", "O", "0"],
            "t": ["t", "T", "7"] }
    try:
        result = table[char.lower()]
    except KeyError:
        result = [char.lower(), char.upper()]
    return result

result = []
# V o l l e y b a l l = 10 chars
for c1 in leet_lookup('v'):
  for c2 in leet_lookup('o'):
    for c3 in leet_lookup('l'):
      for c4 in leet_lookup('l'):
        for c5 in leet_lookup('e'):
          for c6 in leet_lookup('y'):
            for c7 in leet_lookup('b'):
              for c8 in leet_lookup('a'):
                for c9 in leet_lookup('l'):
                  for c10 in leet_lookup('l'):
                    result.append("%s%s%s%s%s%s%s%s%s%s" % (c1, c2, c3, c4, c5, c6, c7, c8, c9, c10))

Solution

  • Once again, it's a problem begging for a solution from itertools.

    For the first problem, simply use product.

    from itertools import product
    for run in product(range(2), repeat=5):
        print(run)
    
    # (0, 0, 0, 0, 0)
    # (0, 0, 0, 0, 1)
    # (0, 0, 0, 1, 0)
    # etc.
    

    If you have more states per indicator, either a) every indicator has the same number of states and you can simply change the argument of range, or b) different indicators have different numbers of states, in which case you could do something similar to the answer below, replacing the leet_lookup call with range and having a list of state numbers instead of lookups.

    For the second one, you probably want to build a list of iterables which you can then use product on.

    lookups = ['v', 'o', 'l', 'l']
    items = [leet_lookup(a) for a in lookups]
    for c in product(*items):
        print(c)
    
    # ('v', 'o', 'l', 'l')
    # ('v', 'o', 'l', 'L')
    # ('v', 'o', 'l', '1')
    # ('v', 'o', 'L', 'l')
    # etc.
    

    In general, if you ever have a problem that involves any sort of combination or transformation of iterators, you should first look for an itertools function or recipe.