How can I improve this heavily nested for-loop?

Question

I have a function which I'd like to optimize, if possible. But I cannot easily tell if there's a better way to refactor (and optimize) this...

Suppose,

keys_in_order = ['A', 'B', 'C', 'D', 'E']
key_table = { 'A': {'A1': 'one', 'A2': 'two', 'A3': 'three', 'A4': 'four'},
'B': {'B1': 'one-one', 'B2': 'two-two', 'B3': 'three-three'},
... # mapping for 'C', 'D' here
'E': {'E1': 'one-one', 'E2': 'two-two', 'E3': 'three-three', 'E6': 'six-six'}
}

The purpose is to feed the above two parameters to the function as below:

def generate_all_possible_key_combinations(keys_in_order, key_table):
    first_key = keys_in_order[0]
    second_key = keys_in_order[1]
    third_key = keys_in_order[2]
    fourth_key = keys_in_order[3]
    fifth_key = keys_in_order[4]

    table_out = [['Demo Group', first_key, second_key, third_key, fourth_key, fifth_key]] # just the header row so that we can write to a CSV file later

    for k1, v1 in key_table[first_key].items():
        for k2, v2 in key_table[second_key].items():
            for k3, v3 in key_table[third_key].items():
                for k4, v4 in key_table[fourth_key].items():
                    for k5, v5 in key_table[fifth_key].items():
                        demo_gp = k1 + k2 + k3 + k4 + k5
                        table_out.append([demo_gp, v1, v2, v3, v4, v5])

    return table_out

so that the goal is to have a table with all possible combination of sub-keys (that is, 'A1B1C1D1E1', 'A1B1C1D1E2', 'A1B1C1D1E3', etc.) along with their corresponding values in key_table .

To me, the current code with five heavily nested loop through the dict key_table is ugly, not to mention it being inefficient computation-wise. Is there a way to improve this? I hope folks from code_review might be able to shed some lights on how I might go about it. Thank you!

Answer 1

I have implemented with an alternative method. Consider as key_table as your main dictionary.

My logic is

1. From this i will get all the possible sub keys from the main dict.

    In [1]: [i.keys() for i in key_table.values()] Out[1]: [['A1', 'A3', 'A2', 'A4'], ['C3', 'C2', 'C1'], ['B1', 'B2', 'B3'], ['E6', 'E1', 'E3', 'E2'], ['D2', 'D3', 'D1']] 

2. Then i made this list of list as a single list.

    In [2]: print [item for sublist in [i.keys() for i in key_table.values()] for item in sublist] ['A1', 'A3', 'A2', 'A4', 'C3', 'C2', 'C1', 'B1', 'B2', 'B3', 'E6', 'E1', 'E3', 'E2', 'D2', 'D3', 'D1'] 

3. With using itertools.combinations implemented the combination of all possible values. It have 5 elements so i given that as a hard code method. You can replace that with len([i.keys() for i in key_table.values()]) if you more values. Here provides an example of itertools.combinations . Then you can understand it.

    In [83]: for i in itertools.combinations(['A1','B1','C1'],2): ....: print i ....: ('A1', 'B1') ('A1', 'C1') ('B1', 'C1') 

Here is the full code with one line implementation.

for item in itertools.combinations([item for sublist in [i.keys() for i in key_table.values()] for item in sublist],5):
    print ''.join(item)

Answer 2

Some optimizations:

• The various key_table[?].items() could be computed before the nested loop
• You could compute partials of demo_gp when they are available: demo_gp12 = k1 + k2 , demo_gp123 = demo_gp12 + k3 , etc. Similar thing could be done with the array of v s.

As @JohnColeman suggested, itertools would be a good place to look to simplifying it.