itertools in python

Reference: doc

In [2]:
import itertools

This module is used to construct iterables in certain way. In general, functions in this module, take one or more exsiting iterables as input, and give a new iterable as output.

accumulate

It is like a iterable version of reduce, instead of giving only a final value, it gives the process list of computation in some sense.

In [4]:
list(itertools.accumulate([11,3,9,7,5],func=min)) # the func kw input is the function used for reduce
Out[4]:
[11, 3, 3, 3, 3]

chain

This function just chain different iterables together

In [6]:
list(itertools.chain([1,2,3],[4,5],[6]))
Out[6]:
[1, 2, 3, 4, 5, 6]
In [89]:
## flatten one level of nested list by chain
list(itertools.chain(*[[1,2],[3,4,[5,6]],[7,8]]))
Out[89]:
[1, 2, 3, 4, [5, 6], 7, 8]

combinations

give combinations of given iterable with given length

In [10]:
list(itertools.combinations([1,2,3,1,2,1],3)) # the result is sorted as the given iterable
Out[10]:
[(1, 2, 3),
 (1, 2, 1),
 (1, 2, 2),
 (1, 2, 1),
 (1, 3, 1),
 (1, 3, 2),
 (1, 3, 1),
 (1, 1, 2),
 (1, 1, 1),
 (1, 2, 1),
 (2, 3, 1),
 (2, 3, 2),
 (2, 3, 1),
 (2, 1, 2),
 (2, 1, 1),
 (2, 2, 1),
 (3, 1, 2),
 (3, 1, 1),
 (3, 2, 1),
 (1, 2, 1)]
In [12]:
list(itertools.combinations("hello world",10))
Out[12]:
[('h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l'),
 ('h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'd'),
 ('h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'l', 'd'),
 ('h', 'e', 'l', 'l', 'o', ' ', 'w', 'r', 'l', 'd'),
 ('h', 'e', 'l', 'l', 'o', ' ', 'o', 'r', 'l', 'd'),
 ('h', 'e', 'l', 'l', 'o', 'w', 'o', 'r', 'l', 'd'),
 ('h', 'e', 'l', 'l', ' ', 'w', 'o', 'r', 'l', 'd'),
 ('h', 'e', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd'),
 ('h', 'e', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd'),
 ('h', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd'),
 ('e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd')]
In [ ]:
# possible implementation
def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    
    pool = tuple(iterable) ## in comment below, suppose the iterable is 0,2,...n-1, same as the indices
    n = len(pool)
    if r > n:
        return
    indices = list(range(r))
    yield tuple(pool[i] for i in indices) ## yield the first combination as 0...r-1
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r: # check against the final combination  n-r ... n-1
                break
        else:  ## if no break, then all combinations have been exausted, return to end the generator
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1 ## indiced after i in a increasing natural order
        yield tuple(pool[i] for i in indices) ## generate this round of combinations based on new indices

combinations with replacement

Similar to combinations, but in the output set, the elements can be duplicated though only one occurence in the input

In [13]:
list(itertools.combinations_with_replacement([1,2,3],5))
Out[13]:
[(1, 1, 1, 1, 1),
 (1, 1, 1, 1, 2),
 (1, 1, 1, 1, 3),
 (1, 1, 1, 2, 2),
 (1, 1, 1, 2, 3),
 (1, 1, 1, 3, 3),
 (1, 1, 2, 2, 2),
 (1, 1, 2, 2, 3),
 (1, 1, 2, 3, 3),
 (1, 1, 3, 3, 3),
 (1, 2, 2, 2, 2),
 (1, 2, 2, 2, 3),
 (1, 2, 2, 3, 3),
 (1, 2, 3, 3, 3),
 (1, 3, 3, 3, 3),
 (2, 2, 2, 2, 2),
 (2, 2, 2, 2, 3),
 (2, 2, 2, 3, 3),
 (2, 2, 3, 3, 3),
 (2, 3, 3, 3, 3),
 (3, 3, 3, 3, 3)]
In [ ]:
## possible implementation
def combinations_with_replacement(iterable, r):
    # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
    pool = tuple(iterable)
    n = len(pool)
    if not n and r: # no iterable, just end
        return
    indices = [0] * r
    yield tuple(pool[i] for i in indices) # yield the first combination 00000
    while True:
        for i in reversed(range(r)):
            if indices[i] != n - 1: # check against the final configuration n-1,n-1...n-1
                break
        else: # if no break, the generator is finished
            return
        indices[i:] = [indices[i] + 1] * (r - i) # update elements after the first non n-1 term plusing one
        yield tuple(pool[i] for i in indices)

compress

something very similar to a mask on the iterable, the second iterable serves as the mask of the first one

In [15]:
list(itertools.compress([1,2,3,4,5],[1,0,1,0,1]))
Out[15]:
[1, 3, 5]

count

count from the start value, with one step more each turn, caution: the returned iterator is infinite

In [3]:
c = itertools.count(10,2)
for i, n in enumerate(c):
    if i>10:
        break
    print(n)
10
12
14
16
18
20
22
24
26
28
30

cycle

generate element according to a given iterator infinite round, Caution: infinite generator

In [4]:
c = itertools.cycle([1,2,3])
for i, n in enumerate(c):
    if i>12:
        break
    print(n)
1
2
3
1
2
3
1
2
3
1
2
3
1

dropwhile

the first argument is the predicate function, when first false evaluation met, generate all element behind regradless of their boolean wrt the first argument. Be caution about the meaning of this function, which is sort of unnatural

In [12]:
list(itertools.dropwhile(lambda x:x<5, [3,5,7,9,1]))
Out[12]:
[5, 7, 9, 1]
In [13]:
list(itertools.dropwhile(lambda x:x>5, [3,5,7,9,1]))
Out[13]:
[3, 5, 7, 9, 1]

filterfalse

return all elements evaluated to false, contrary to the builtin filter function

In [11]:
list(itertools.filterfalse(lambda x:x<5, [3,5,7,9,1]))
Out[11]:
[5, 7, 9]

groupby

return iterator of iterators, each group of iterator is group by the key function with same return value, the input data should already be sorted by the key, because the implementation is to detect the change of the key value and then yield

In [32]:
gb = list(itertools.groupby([(1,2),(2,2),(2,3),(3,3),(1,3)],key= lambda x:x[1]))
gb
Out[32]:
[(2, <itertools._grouper at 0x10d0ed7b8>),
 (3, <itertools._grouper at 0x10d0ed908>)]
In [29]:
## right way
l = []
for i in itertools.groupby([(1,2),(2,2),(2,3),(3,3),(1,3)],key= lambda x:x[1]):
    l.append(list(i[1]))
In [30]:
l
Out[30]:
[[(1, 2), (2, 2)], [(2, 3), (3, 3), (1, 3)]]

Caution: The returned group is itself an iterator that shares the underlying iterable with groupby(). Because the source is shared, when the groupby() object is advanced, the previous group is no longer visible. So, if that data is needed later, it should be stored as a list:

In [34]:
# wrongway
list(gb[0][1]) ## the other elements are gone since it is a generator
Out[34]:
[]

islice

In [36]:
list(itertools.islice(range(20),2,8,2)) == list(range(20))[2:8:2] # one line is the whole story of islice
Out[36]:
True

permutations

order matter version of combination

In [37]:
list(itertools.permutations("abcd",2))
Out[37]:
[('a', 'b'),
 ('a', 'c'),
 ('a', 'd'),
 ('b', 'a'),
 ('b', 'c'),
 ('b', 'd'),
 ('c', 'a'),
 ('c', 'b'),
 ('c', 'd'),
 ('d', 'a'),
 ('d', 'b'),
 ('d', 'c')]
In [39]:
list(itertools.permutations([1,2,3,2],2))
Out[39]:
[(1, 2),
 (1, 3),
 (1, 2),
 (2, 1),
 (2, 3),
 (2, 2),
 (3, 1),
 (3, 2),
 (3, 2),
 (2, 1),
 (2, 2),
 (2, 3)]
In [100]:
## possible implementation
def permutations(iterable, r=None):
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
    # permutations(range(3)) --> 012 021 102 120 201 210
    pool = tuple(iterable)
    n = len(pool)
    r = n if r is None else r
    if r > n:
        return
    indices = list(range(n)) ## the first one is 0,1...r-1
    cycles = list(range(n, n-r, -1)) ## cycle = n-1, n-2,... n-r+1 
    ## the ith pos of cycle is the next larger ele pos after indices[i], which is indices[-j], and plus 1
    ## if cycle[i]=0 after the -1 process, it means there is nothing larger than indices[i] after it
    print(f"initial indices: {indices}")
    print(f"initial cycles: {cycles}")
    yield tuple(pool[i] for i in indices[:r])
    while n:
        print("-------new round of yield-------")
        for i in reversed(range(r)):
            print(f"i={i}")
            cycles[i] -= 1
            if cycles[i] == 0:
                print(f"update cycles, since cycle[{i}] = 0")
                indices[i:] = indices[i+1:] + indices[i:i+1]  # push the ith pos of indices to the back,
                # come back to the initial pos of this cycle, still in increasing order after i
                cycles[i] = n - i # reset the i pos of cycle to the original value, since the state of indices  
                # after ith pos is coming back to original increasing state
                print(f"cycles: {cycles}")
                print(f"indices: {indices}")
            else: # cycle the i pos with later pos to make sure every element occur within the first r window
                j = cycles[i] 
                print(f"iterchange {i}, {-j} in indices since cycle[{i}]={j}") # interchange i and the next larger one
                indices[i], indices[-j] = indices[-j], indices[i] # the real cycle process
                print(f"cycles: {cycles}")
                print(f"yield indices: {indices}")
                yield tuple(pool[i] for i in indices[:r])
                print("-------break this round of yield---------")
                break
        else: # if no break, end
            return
In [104]:
list(permutations([0,1,2,3,4],3)) ## see the debug info to learn about how this implementation works
initial indices: [0, 1, 2, 3, 4]
initial cycles: [5, 4, 3]
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [5, 4, 2]
yield indices: [0, 1, 3, 2, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [5, 4, 1]
yield indices: [0, 1, 4, 2, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [5, 4, 3]
indices: [0, 1, 2, 3, 4]
i=1
iterchange 1, -3 in indices since cycle[1]=3
cycles: [5, 3, 3]
yield indices: [0, 2, 1, 3, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [5, 3, 2]
yield indices: [0, 2, 3, 1, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [5, 3, 1]
yield indices: [0, 2, 4, 1, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [5, 3, 3]
indices: [0, 2, 1, 3, 4]
i=1
iterchange 1, -2 in indices since cycle[1]=2
cycles: [5, 2, 3]
yield indices: [0, 3, 1, 2, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [5, 2, 2]
yield indices: [0, 3, 2, 1, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [5, 2, 1]
yield indices: [0, 3, 4, 1, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [5, 2, 3]
indices: [0, 3, 1, 2, 4]
i=1
iterchange 1, -1 in indices since cycle[1]=1
cycles: [5, 1, 3]
yield indices: [0, 4, 1, 2, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [5, 1, 2]
yield indices: [0, 4, 2, 1, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [5, 1, 1]
yield indices: [0, 4, 3, 1, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [5, 1, 3]
indices: [0, 4, 1, 2, 3]
i=1
update cycles, since cycle[1] = 0
cycles: [5, 4, 3]
indices: [0, 1, 2, 3, 4]
i=0
iterchange 0, -4 in indices since cycle[0]=4
cycles: [4, 4, 3]
yield indices: [1, 0, 2, 3, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [4, 4, 2]
yield indices: [1, 0, 3, 2, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [4, 4, 1]
yield indices: [1, 0, 4, 2, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [4, 4, 3]
indices: [1, 0, 2, 3, 4]
i=1
iterchange 1, -3 in indices since cycle[1]=3
cycles: [4, 3, 3]
yield indices: [1, 2, 0, 3, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [4, 3, 2]
yield indices: [1, 2, 3, 0, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [4, 3, 1]
yield indices: [1, 2, 4, 0, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [4, 3, 3]
indices: [1, 2, 0, 3, 4]
i=1
iterchange 1, -2 in indices since cycle[1]=2
cycles: [4, 2, 3]
yield indices: [1, 3, 0, 2, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [4, 2, 2]
yield indices: [1, 3, 2, 0, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [4, 2, 1]
yield indices: [1, 3, 4, 0, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [4, 2, 3]
indices: [1, 3, 0, 2, 4]
i=1
iterchange 1, -1 in indices since cycle[1]=1
cycles: [4, 1, 3]
yield indices: [1, 4, 0, 2, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [4, 1, 2]
yield indices: [1, 4, 2, 0, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [4, 1, 1]
yield indices: [1, 4, 3, 0, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [4, 1, 3]
indices: [1, 4, 0, 2, 3]
i=1
update cycles, since cycle[1] = 0
cycles: [4, 4, 3]
indices: [1, 0, 2, 3, 4]
i=0
iterchange 0, -3 in indices since cycle[0]=3
cycles: [3, 4, 3]
yield indices: [2, 0, 1, 3, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [3, 4, 2]
yield indices: [2, 0, 3, 1, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [3, 4, 1]
yield indices: [2, 0, 4, 1, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [3, 4, 3]
indices: [2, 0, 1, 3, 4]
i=1
iterchange 1, -3 in indices since cycle[1]=3
cycles: [3, 3, 3]
yield indices: [2, 1, 0, 3, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [3, 3, 2]
yield indices: [2, 1, 3, 0, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [3, 3, 1]
yield indices: [2, 1, 4, 0, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [3, 3, 3]
indices: [2, 1, 0, 3, 4]
i=1
iterchange 1, -2 in indices since cycle[1]=2
cycles: [3, 2, 3]
yield indices: [2, 3, 0, 1, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [3, 2, 2]
yield indices: [2, 3, 1, 0, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [3, 2, 1]
yield indices: [2, 3, 4, 0, 1]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [3, 2, 3]
indices: [2, 3, 0, 1, 4]
i=1
iterchange 1, -1 in indices since cycle[1]=1
cycles: [3, 1, 3]
yield indices: [2, 4, 0, 1, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [3, 1, 2]
yield indices: [2, 4, 1, 0, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [3, 1, 1]
yield indices: [2, 4, 3, 0, 1]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [3, 1, 3]
indices: [2, 4, 0, 1, 3]
i=1
update cycles, since cycle[1] = 0
cycles: [3, 4, 3]
indices: [2, 0, 1, 3, 4]
i=0
iterchange 0, -2 in indices since cycle[0]=2
cycles: [2, 4, 3]
yield indices: [3, 0, 1, 2, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [2, 4, 2]
yield indices: [3, 0, 2, 1, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [2, 4, 1]
yield indices: [3, 0, 4, 1, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [2, 4, 3]
indices: [3, 0, 1, 2, 4]
i=1
iterchange 1, -3 in indices since cycle[1]=3
cycles: [2, 3, 3]
yield indices: [3, 1, 0, 2, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [2, 3, 2]
yield indices: [3, 1, 2, 0, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [2, 3, 1]
yield indices: [3, 1, 4, 0, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [2, 3, 3]
indices: [3, 1, 0, 2, 4]
i=1
iterchange 1, -2 in indices since cycle[1]=2
cycles: [2, 2, 3]
yield indices: [3, 2, 0, 1, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [2, 2, 2]
yield indices: [3, 2, 1, 0, 4]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [2, 2, 1]
yield indices: [3, 2, 4, 0, 1]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [2, 2, 3]
indices: [3, 2, 0, 1, 4]
i=1
iterchange 1, -1 in indices since cycle[1]=1
cycles: [2, 1, 3]
yield indices: [3, 4, 0, 1, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [2, 1, 2]
yield indices: [3, 4, 1, 0, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [2, 1, 1]
yield indices: [3, 4, 2, 0, 1]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [2, 1, 3]
indices: [3, 4, 0, 1, 2]
i=1
update cycles, since cycle[1] = 0
cycles: [2, 4, 3]
indices: [3, 0, 1, 2, 4]
i=0
iterchange 0, -1 in indices since cycle[0]=1
cycles: [1, 4, 3]
yield indices: [4, 0, 1, 2, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [1, 4, 2]
yield indices: [4, 0, 2, 1, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [1, 4, 1]
yield indices: [4, 0, 3, 1, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [1, 4, 3]
indices: [4, 0, 1, 2, 3]
i=1
iterchange 1, -3 in indices since cycle[1]=3
cycles: [1, 3, 3]
yield indices: [4, 1, 0, 2, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [1, 3, 2]
yield indices: [4, 1, 2, 0, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [1, 3, 1]
yield indices: [4, 1, 3, 0, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [1, 3, 3]
indices: [4, 1, 0, 2, 3]
i=1
iterchange 1, -2 in indices since cycle[1]=2
cycles: [1, 2, 3]
yield indices: [4, 2, 0, 1, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [1, 2, 2]
yield indices: [4, 2, 1, 0, 3]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [1, 2, 1]
yield indices: [4, 2, 3, 0, 1]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [1, 2, 3]
indices: [4, 2, 0, 1, 3]
i=1
iterchange 1, -1 in indices since cycle[1]=1
cycles: [1, 1, 3]
yield indices: [4, 3, 0, 1, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -2 in indices since cycle[2]=2
cycles: [1, 1, 2]
yield indices: [4, 3, 1, 0, 2]
-------break this round of yield---------
-------new round of yield-------
i=2
iterchange 2, -1 in indices since cycle[2]=1
cycles: [1, 1, 1]
yield indices: [4, 3, 2, 0, 1]
-------break this round of yield---------
-------new round of yield-------
i=2
update cycles, since cycle[2] = 0
cycles: [1, 1, 3]
indices: [4, 3, 0, 1, 2]
i=1
update cycles, since cycle[1] = 0
cycles: [1, 4, 3]
indices: [4, 0, 1, 2, 3]
i=0
update cycles, since cycle[0] = 0
cycles: [5, 4, 3]
indices: [0, 1, 2, 3, 4]
Out[104]:
[(0, 1, 2),
 (0, 1, 3),
 (0, 1, 4),
 (0, 2, 1),
 (0, 2, 3),
 (0, 2, 4),
 (0, 3, 1),
 (0, 3, 2),
 (0, 3, 4),
 (0, 4, 1),
 (0, 4, 2),
 (0, 4, 3),
 (1, 0, 2),
 (1, 0, 3),
 (1, 0, 4),
 (1, 2, 0),
 (1, 2, 3),
 (1, 2, 4),
 (1, 3, 0),
 (1, 3, 2),
 (1, 3, 4),
 (1, 4, 0),
 (1, 4, 2),
 (1, 4, 3),
 (2, 0, 1),
 (2, 0, 3),
 (2, 0, 4),
 (2, 1, 0),
 (2, 1, 3),
 (2, 1, 4),
 (2, 3, 0),
 (2, 3, 1),
 (2, 3, 4),
 (2, 4, 0),
 (2, 4, 1),
 (2, 4, 3),
 (3, 0, 1),
 (3, 0, 2),
 (3, 0, 4),
 (3, 1, 0),
 (3, 1, 2),
 (3, 1, 4),
 (3, 2, 0),
 (3, 2, 1),
 (3, 2, 4),
 (3, 4, 0),
 (3, 4, 1),
 (3, 4, 2),
 (4, 0, 1),
 (4, 0, 2),
 (4, 0, 3),
 (4, 1, 0),
 (4, 1, 2),
 (4, 1, 3),
 (4, 2, 0),
 (4, 2, 1),
 (4, 2, 3),
 (4, 3, 0),
 (4, 3, 1),
 (4, 3, 2)]

product

Cartesian product of the input iterator, multiple iterators can be taken as input

In [43]:
list(itertools.product([1,2,3,4],"hi"))
Out[43]:
[(1, 'h'),
 (1, 'i'),
 (2, 'h'),
 (2, 'i'),
 (3, 'h'),
 (3, 'i'),
 (4, 'h'),
 (4, 'i')]
In [42]:
list(itertools.product([1,2,3,4],[5,6,7,8],repeat=2))
Out[42]:
[(1, 5, 1, 5),
 (1, 5, 1, 6),
 (1, 5, 1, 7),
 (1, 5, 1, 8),
 (1, 5, 2, 5),
 (1, 5, 2, 6),
 (1, 5, 2, 7),
 (1, 5, 2, 8),
 (1, 5, 3, 5),
 (1, 5, 3, 6),
 (1, 5, 3, 7),
 (1, 5, 3, 8),
 (1, 5, 4, 5),
 (1, 5, 4, 6),
 (1, 5, 4, 7),
 (1, 5, 4, 8),
 (1, 6, 1, 5),
 (1, 6, 1, 6),
 (1, 6, 1, 7),
 (1, 6, 1, 8),
 (1, 6, 2, 5),
 (1, 6, 2, 6),
 (1, 6, 2, 7),
 (1, 6, 2, 8),
 (1, 6, 3, 5),
 (1, 6, 3, 6),
 (1, 6, 3, 7),
 (1, 6, 3, 8),
 (1, 6, 4, 5),
 (1, 6, 4, 6),
 (1, 6, 4, 7),
 (1, 6, 4, 8),
 (1, 7, 1, 5),
 (1, 7, 1, 6),
 (1, 7, 1, 7),
 (1, 7, 1, 8),
 (1, 7, 2, 5),
 (1, 7, 2, 6),
 (1, 7, 2, 7),
 (1, 7, 2, 8),
 (1, 7, 3, 5),
 (1, 7, 3, 6),
 (1, 7, 3, 7),
 (1, 7, 3, 8),
 (1, 7, 4, 5),
 (1, 7, 4, 6),
 (1, 7, 4, 7),
 (1, 7, 4, 8),
 (1, 8, 1, 5),
 (1, 8, 1, 6),
 (1, 8, 1, 7),
 (1, 8, 1, 8),
 (1, 8, 2, 5),
 (1, 8, 2, 6),
 (1, 8, 2, 7),
 (1, 8, 2, 8),
 (1, 8, 3, 5),
 (1, 8, 3, 6),
 (1, 8, 3, 7),
 (1, 8, 3, 8),
 (1, 8, 4, 5),
 (1, 8, 4, 6),
 (1, 8, 4, 7),
 (1, 8, 4, 8),
 (2, 5, 1, 5),
 (2, 5, 1, 6),
 (2, 5, 1, 7),
 (2, 5, 1, 8),
 (2, 5, 2, 5),
 (2, 5, 2, 6),
 (2, 5, 2, 7),
 (2, 5, 2, 8),
 (2, 5, 3, 5),
 (2, 5, 3, 6),
 (2, 5, 3, 7),
 (2, 5, 3, 8),
 (2, 5, 4, 5),
 (2, 5, 4, 6),
 (2, 5, 4, 7),
 (2, 5, 4, 8),
 (2, 6, 1, 5),
 (2, 6, 1, 6),
 (2, 6, 1, 7),
 (2, 6, 1, 8),
 (2, 6, 2, 5),
 (2, 6, 2, 6),
 (2, 6, 2, 7),
 (2, 6, 2, 8),
 (2, 6, 3, 5),
 (2, 6, 3, 6),
 (2, 6, 3, 7),
 (2, 6, 3, 8),
 (2, 6, 4, 5),
 (2, 6, 4, 6),
 (2, 6, 4, 7),
 (2, 6, 4, 8),
 (2, 7, 1, 5),
 (2, 7, 1, 6),
 (2, 7, 1, 7),
 (2, 7, 1, 8),
 (2, 7, 2, 5),
 (2, 7, 2, 6),
 (2, 7, 2, 7),
 (2, 7, 2, 8),
 (2, 7, 3, 5),
 (2, 7, 3, 6),
 (2, 7, 3, 7),
 (2, 7, 3, 8),
 (2, 7, 4, 5),
 (2, 7, 4, 6),
 (2, 7, 4, 7),
 (2, 7, 4, 8),
 (2, 8, 1, 5),
 (2, 8, 1, 6),
 (2, 8, 1, 7),
 (2, 8, 1, 8),
 (2, 8, 2, 5),
 (2, 8, 2, 6),
 (2, 8, 2, 7),
 (2, 8, 2, 8),
 (2, 8, 3, 5),
 (2, 8, 3, 6),
 (2, 8, 3, 7),
 (2, 8, 3, 8),
 (2, 8, 4, 5),
 (2, 8, 4, 6),
 (2, 8, 4, 7),
 (2, 8, 4, 8),
 (3, 5, 1, 5),
 (3, 5, 1, 6),
 (3, 5, 1, 7),
 (3, 5, 1, 8),
 (3, 5, 2, 5),
 (3, 5, 2, 6),
 (3, 5, 2, 7),
 (3, 5, 2, 8),
 (3, 5, 3, 5),
 (3, 5, 3, 6),
 (3, 5, 3, 7),
 (3, 5, 3, 8),
 (3, 5, 4, 5),
 (3, 5, 4, 6),
 (3, 5, 4, 7),
 (3, 5, 4, 8),
 (3, 6, 1, 5),
 (3, 6, 1, 6),
 (3, 6, 1, 7),
 (3, 6, 1, 8),
 (3, 6, 2, 5),
 (3, 6, 2, 6),
 (3, 6, 2, 7),
 (3, 6, 2, 8),
 (3, 6, 3, 5),
 (3, 6, 3, 6),
 (3, 6, 3, 7),
 (3, 6, 3, 8),
 (3, 6, 4, 5),
 (3, 6, 4, 6),
 (3, 6, 4, 7),
 (3, 6, 4, 8),
 (3, 7, 1, 5),
 (3, 7, 1, 6),
 (3, 7, 1, 7),
 (3, 7, 1, 8),
 (3, 7, 2, 5),
 (3, 7, 2, 6),
 (3, 7, 2, 7),
 (3, 7, 2, 8),
 (3, 7, 3, 5),
 (3, 7, 3, 6),
 (3, 7, 3, 7),
 (3, 7, 3, 8),
 (3, 7, 4, 5),
 (3, 7, 4, 6),
 (3, 7, 4, 7),
 (3, 7, 4, 8),
 (3, 8, 1, 5),
 (3, 8, 1, 6),
 (3, 8, 1, 7),
 (3, 8, 1, 8),
 (3, 8, 2, 5),
 (3, 8, 2, 6),
 (3, 8, 2, 7),
 (3, 8, 2, 8),
 (3, 8, 3, 5),
 (3, 8, 3, 6),
 (3, 8, 3, 7),
 (3, 8, 3, 8),
 (3, 8, 4, 5),
 (3, 8, 4, 6),
 (3, 8, 4, 7),
 (3, 8, 4, 8),
 (4, 5, 1, 5),
 (4, 5, 1, 6),
 (4, 5, 1, 7),
 (4, 5, 1, 8),
 (4, 5, 2, 5),
 (4, 5, 2, 6),
 (4, 5, 2, 7),
 (4, 5, 2, 8),
 (4, 5, 3, 5),
 (4, 5, 3, 6),
 (4, 5, 3, 7),
 (4, 5, 3, 8),
 (4, 5, 4, 5),
 (4, 5, 4, 6),
 (4, 5, 4, 7),
 (4, 5, 4, 8),
 (4, 6, 1, 5),
 (4, 6, 1, 6),
 (4, 6, 1, 7),
 (4, 6, 1, 8),
 (4, 6, 2, 5),
 (4, 6, 2, 6),
 (4, 6, 2, 7),
 (4, 6, 2, 8),
 (4, 6, 3, 5),
 (4, 6, 3, 6),
 (4, 6, 3, 7),
 (4, 6, 3, 8),
 (4, 6, 4, 5),
 (4, 6, 4, 6),
 (4, 6, 4, 7),
 (4, 6, 4, 8),
 (4, 7, 1, 5),
 (4, 7, 1, 6),
 (4, 7, 1, 7),
 (4, 7, 1, 8),
 (4, 7, 2, 5),
 (4, 7, 2, 6),
 (4, 7, 2, 7),
 (4, 7, 2, 8),
 (4, 7, 3, 5),
 (4, 7, 3, 6),
 (4, 7, 3, 7),
 (4, 7, 3, 8),
 (4, 7, 4, 5),
 (4, 7, 4, 6),
 (4, 7, 4, 7),
 (4, 7, 4, 8),
 (4, 8, 1, 5),
 (4, 8, 1, 6),
 (4, 8, 1, 7),
 (4, 8, 1, 8),
 (4, 8, 2, 5),
 (4, 8, 2, 6),
 (4, 8, 2, 7),
 (4, 8, 2, 8),
 (4, 8, 3, 5),
 (4, 8, 3, 6),
 (4, 8, 3, 7),
 (4, 8, 3, 8),
 (4, 8, 4, 5),
 (4, 8, 4, 6),
 (4, 8, 4, 7),
 (4, 8, 4, 8)]
In [ ]:
## possible implementation
def product(*args, repeat=1):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    pools = [tuple(pool) for pool in args] * repeat # pools = [(A,B,C,D),(x,y), [A,B,C,D],[x,y]]
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)

repeat

It generates one thing forever, unless time argument is specified, caution: possible infinite generator

In [64]:
list(itertools.repeat(1, times=10))
Out[64]:
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
In [66]:
list(map(pow, range(10), itertools.repeat(2)))
Out[66]:
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
In [67]:
## a better map approach avoid repeat
list(map(lambda x: pow(x,2), range(10)))
Out[67]:
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

startmap

In [68]:
list(itertools.starmap(pow, [(2,5), (3,2), (10,3)]))
Out[68]:
[32, 9, 1000]
In [ ]:
## simple implementation
def starmap(function, iterable):
    # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
    for args in iterable:
        yield function(*args)

takewhile

It generates elements from the second input as long as the first input func is evaluated to be true. Namely stop generator when the first false is met. Caution: make sure the meaning of this function, it is not so natural.

In [74]:
list(itertools.takewhile(lambda x: x<5, [1,6,9,1]))
Out[74]:
[1]

tee

This turns one iterator into multiple copies

In [80]:
ls = list(itertools.tee([1,2,3,4],5))
In [81]:
list(ls[0]),list(ls[1])
Out[81]:
([1, 2, 3, 4], [1, 2, 3, 4])
In [82]:
ls
Out[82]:
[<itertools._tee at 0x10d1d49c8>,
 <itertools._tee at 0x10d1c5208>,
 <itertools._tee at 0x10d1c5188>,
 <itertools._tee at 0x10d1c4648>,
 <itertools._tee at 0x10d1c46c8>]

zip_longest

A more robust and general version of zip. The iterables can be in unequal length. The missing value is treated as fillvalue flag in the function.

In [84]:
list(itertools.zip_longest([1,2,3,4],[5,6,7,8]))
Out[84]:
[(1, 5), (2, 6), (3, 7), (4, 8)]
In [85]:
list(itertools.zip_longest([1,2,3,4],[5,6,7,8,9]))
Out[85]:
[(1, 5), (2, 6), (3, 7), (4, 8), (None, 9)]
In [86]:
list(itertools.zip_longest([1,2,3,4],[5,6,7,8,9,10],fillvalue='*'))
Out[86]:
[(1, 5), (2, 6), (3, 7), (4, 8), ('*', 9), ('*', 10)]
In [87]:
list(itertools.zip_longest([1,2,3,4],[5,6,7,8,9,10],[11,12],fillvalue=0))
Out[87]:
[(1, 5, 11), (2, 6, 12), (3, 7, 0), (4, 8, 0), (0, 9, 0), (0, 10, 0)]