Welcome, guest | Sign In | My Account | Store | Cart

Another generator for an arbitrary number of 'for' loops (Python recipe) by Paddy McCarthy
ActiveState Code (http://code.activestate.com/recipes/502199/)

Saw recipe 502194 and wondered if I could do similar. I initially came up with function comb that enumerates all combinations. I wanted a generator however and so abandoned that approach for comb2 which is a little more complex.

Python, 79 lines
Download
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def comb(*sequences):
    '''
    combinations of multiple sequences so you don't have
    to write nested for loops
    
    >>> from pprint import pprint as pp
    >>> pp(comb(['Guido','Larry'], ['knows','loves'], ['Phyton','Purl']))
    [['Guido', 'knows', 'Phyton'],
     ['Guido', 'knows', 'Purl'],
     ['Guido', 'loves', 'Phyton'],
     ['Guido', 'loves', 'Purl'],
     ['Larry', 'knows', 'Phyton'],
     ['Larry', 'knows', 'Purl'],
     ['Larry', 'loves', 'Phyton'],
     ['Larry', 'loves', 'Purl']]
    >>> 
    '''
    combinations = [[seq] for seq in sequences[0]]
    for seq in sequences[1:]:
        combinations = [comb+[item]
                        for comb in combinations
                        for item in seq ]
    return combinations

def comb2(*sequences):
    '''
    Generator of combinations of multiple sequences so you don't have
    to write nested for loops

    Note: rightmost sequence changes the quickest as if it were the inner loop.    
    
    >>> for x in comb2(['Guido','Larry'], ['knows','loves','hates'], ['Phyton','Purl','Wuby','Auk']):
    ... 	print x
    ... 	
    ['Guido', 'knows', 'Phyton']
    ['Guido', 'knows', 'Purl']
    ['Guido', 'knows', 'Wuby']
    ['Guido', 'knows', 'Auk']
    ['Guido', 'loves', 'Phyton']
    ['Guido', 'loves', 'Purl']
    ['Guido', 'loves', 'Wuby']
    ['Guido', 'loves', 'Auk']
    ['Guido', 'hates', 'Phyton']
    ['Guido', 'hates', 'Purl']
    ['Guido', 'hates', 'Wuby']
    ['Guido', 'hates', 'Auk']
    ['Larry', 'knows', 'Phyton']
    ['Larry', 'knows', 'Purl']
    ['Larry', 'knows', 'Wuby']
    ['Larry', 'knows', 'Auk']
    ['Larry', 'loves', 'Phyton']
    ['Larry', 'loves', 'Purl']
    ['Larry', 'loves', 'Wuby']
    ['Larry', 'loves', 'Auk']
    ['Larry', 'hates', 'Phyton']
    ['Larry', 'hates', 'Purl']
    ['Larry', 'hates', 'Wuby']
    ['Larry', 'hates', 'Auk']
    >>>

    The algorithm relies on noting the following and generalising:

    If you had three sequences of length 2, 3, and 4 left-to-right,
    Then the indices of the *elements* for all combinations can be
    generated with:
        for x in range(2*3*4): 
            print ( (x/(3*4*1))%2, (x/(4*1))%3, (x/(1))%4 )
    
    '''
    import operator
    
    lengths = [len(seq) for seq in sequences]
    range_len_seq = range(len(sequences))
    max_count = reduce(operator.mul, lengths)
    _tmp = lengths + [1]        # append multiplicative identity
    dividers = [reduce(operator.mul, _tmp[-x-1:]) for x in range_len_seq][::-1]
    modulos = lengths
    for n in range(max_count):
        yield [sequences[r][(n/dividers[r])%modulos[r]] for r in range_len_seq]

Deeply nested code is usually code in need of refactoring, but if you still think you need it then comp2 will save you some levels of indentation.

  • Donald 'Paddy' McCarthy.
Tags: algorithms