pythonのitertoolsのrecipeのメモ

公式にitertoolsのrecipe集があるのでそれのメモ．
バージョン: python3.3
参考: itertools — Functions creating iterators for efficient looping — Python 3.7.3 documentation

・take(n,iterable)
最初のn個の要素をリストとして返します

from itertools import islice 
def take(n, iterable):
    return list(islice(iterable, n))

例

In [76]: take(3,[1,2,3,4,5])
Out[76]: [1, 2, 3]

・tabulate(function,start=0)
function(0),function(1),... を返すイテレータを返します．

from itertools import count 
def tabulate(function, start=0):
    return map(function, count(start))

例

In [82]: i = tabulate(lambda x: x*x)

In [83]: next(i)
Out[83]: 0

In [84]: next(i)
Out[84]: 1

In [85]: next(i)
Out[85]: 4

In [86]: next(i)
Out[86]: 9

・consume(iterator,n)
イテレータをn進める，もしnがNoneならイテレータを末尾まで進める

import collections
from itertools import islice
def consume(iterator, n):
    if n is None:
        collections.deque(iterator, maxlen=0)
    else:
        next(islice(iterator, n, n), None)

例

In [105]: it = iter([1,2,3,4,5,6,7,8,9,10])

In [106]: consume(it,n=5)

In [107]: for i in it:
   .....:     print(i)
   .....:
6
7
8
9
10

・nth(iterable,n,default=None)
n番目の要素を返します．もしnが範囲外ならdefaultを返します

from itertools import islice
def nth(iterable, n, default=None):
    return next(islice(iterable, n, None), default)

例

In [110]: nth([1,2,3],2)
Out[110]: 3

・quantify(iterable,pred=bool)
predがTrueになる要素の数を返します

def quantify(iterable, pred=bool):
    return sum(map(pred, iterable))

例

In [115]: quantify([1,2,3,4,5],pred=lambda x: x<3)
Out[115]: 2  # 3より小さい数の数

・padnone(iterable):
iterableな要素の後にNoneを付け加えたイテレータを返す

from itertools imprt chain,repeat
def padnone(iterable):
    return chain(iterable, repeat(None))

例

In [125]: a = padnone([1,2,3])

In [126]: next(a)
Out[126]: 1

In [127]: next(a)
Out[127]: 2

In [128]: next(a)
Out[128]: 3

In [129]: next(a)

・ncycles(iterable,n):
iterableをn回繰り返すイテレータを返す

frmo itertools import chain,repeat
def ncycles(iterable, n):
    return chain.from_iterable(repeat(tuple(iterable), n))

例

In [136]: it = ncycles([1,2],3)

In [137]: for i in it:
   .....:     print(i)
   .....:
1
2
1
2
1
2

・dotproduct(vec1,vec2)
vec1,vec2の内積を返します．

import operator
def dotproduct(vec1, vec2):
    return sum(map(operator.mul, vec1, vec2))

例

In [141]: dotproduct([1,2],[3,4])
Out[141]: 11

・flatten(listOfLists)
リストのリストを一つのリストにまとめます

from itertools import chain
def flatten(listOfLists):
    return chain.from_iterable(listOfLists)

例

In [144]: it = flatten([[1,2],['a','b']])

In [145]: for i in it:
   .....:     print(i)
   .....:
1
2
a
b

・repeatfunc(func,times=None,*args)
funcをtimes回リピートした結果を返すイテレータを返します．
timesを指定しない場合は無限リストになります．

from itertools import starmap,repeat
def repeatfunc(func, times=None, *args):
    if times is None:
        return starmap(func, repeat(args))
    return starmap(func, repeat(args, times))

例

In [150]: r = repeatfunc(random.random)

In [151]: next(r)
Out[151]: 0.5425125491886116

In [152]: next(r)
Out[152]: 0.37101369831757924

・pairwise(iterable)
1つだけ指す場所が違う2つのイテレータを返します．連続した2つの要素を操作するのに使えます．

from itertools import tee
def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

例

In [163]: xs = [1,4,9,16,25]

# xs[i] - xs[i-1] を計算したい
In [164]: for a,b in pairwise(xs):
   .....:     print(b-a)
   .....:
3
5
7
9

・grouper(n,iterable,fillvalue=None)
要素をn個ずつにまとめます．要素数がnで割り切れなかったらfillvalueで埋められます．

from itertools import zip_longest
def grouper(n, iterable, fillvalue=None):
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

例

In [168]: it = grouper(3,'ABCDEFG','x')

In [169]: for i in it:
   .....:     print(i)
   .....:
('A', 'B', 'C')
('D', 'E', 'F')
('G', 'x', 'x')

なんでこういう動作をするかというと，args = [iter(iterable)]*n とイテレータを複製していますが，このイテレータは全て同一であるため，いずれか一つのイテレータを操作すると他のイテレータも操作されるためです．

・roundrobin(*iterables)
それぞれのiterableなものから順番に要素をとってきます．

from itertools import cycle,islice
def roundrobin(*iterables):
    # Recipe credited to George Sakkis
    pending = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while pending:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            pending -= 1
            nexts = cycle(islice(nexts, pending))

例

In [171]: it = roundrobin([1,2],'AB','xy')

In [172]: for i in it:
   .....:     print(i)
   .....:
1
A
x
2
B
y

・partition(pred,iterable)
predがTrueかどうかに応じてiterableを2つに分割します

from itertools import filterfalse
def partition(pred, iterable):
    t1, t2 = tee(iterable)
    return filterfalse(pred, t1), filter(pred, t2)

例

In [174]: it1,it2 = partition(lambda x: x<3, [1,2,3,4,5])

In [175]: for i in it1:
    print(i)
   .....:
3
4
5

In [176]: for i in it2:
    print(i)
   .....:
1
2

・powerset(iterable):
iterableの冪集合を返します

from itertools import chain,combinations
def powerset(iterable):
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))

例

In [178]: it = powerset([1,2,3])

In [179]: for i in it:
   .....:     print(i)
   .....:
()
(1,)
(2,)
(3,)
(1, 2)
(1, 3)
(2, 3)
(1, 2, 3)

・unique_everseen(iterable,key=None)
iterableな要素のうちユニークなもののみを返すイテレータを返します

from itertools import filterfalse
def unique_everseen(iterable, key=None):
    seen = set()
    seen_add = seen.add
    if key is None:
        for element in filterfalse(seen.__contains__, iterable):
            seen_add(element)
            yield element
    else:
        for element in iterable:
            k = key(element)
            if k not in seen:
                seen_add(k)
                yield element

例

In [189]: it = unique_everseen([1,2,1,1,2,4,2,5])

In [190]: for i in it:
    print(i)
   .....:
1
2
4
5

・unique_justseen(iterable,key=None)
連続する要素を1つにまとめます．

from operator import itemgetter
from itertools import groupby
def unique_justseen(iterable, key=None):
    return map(next, map(itemgetter(1), groupby(iterable, key)))

例

In [200]: it = unique_justseen([1,1,2,3,3,1,2,2])

In [201]: for i in it:
   .....:     print(i)
   .....:
1
2
3
1
2

下の例でも結果は同じだけど上のようにやると遅延評価できる

In [206]: def unique_justseen(iterable,key=None):
    return map(itemgetter(0),itertools.groupby(iterable,key))
   .....:

In [207]: for i in it:
    print(i)
   .....:

In [208]: it = unique_justseen([1,1,2,3,3,1,2,2])

In [209]: for i in it:
    print(i)
   .....:
1
2
3
1
2

・iter_except(func,exception,first=None)
funcを例外が発生するまで呼び続けます．firstは最初に1回だけ実行される(dbの初期化とかに使う)．

def iter_except(func, exception, first=None):
    try:
        if first is not None:
            yield first()
        while 1:
            yield func()
    except exception:
        pass

例

In [219]: it = iter_except(d.popitem,KeyError)

In [220]: for i in it:       # KeyErrorで中断されることなく処理できる
   .....:     print(i)
   .....:
('a', 1)
('b', 2)
('c', 3)

・random_product(*args,repeat=1)
直積をランダムに1つ返します

import random
def random_product(*args, repeat=1):
    pools = [tuple(pool) for pool in args] * repeat
    return tuple(random.choice(pool) for pool in pools)

例

In [277]: random_product([1,2],['x','y'])
Out[277]: (1, 'x')

In [278]: random_product([1,2],['x','y'])
Out[278]: (1, 'y')

・random_permutation(iterable,r=None)
順列をランダムに1つ返します

import random
def random_permutation(iterable, r=None):
    pool = tuple(iterable)
    r = len(pool) if r is None else r
    return tuple(random.sample(pool, r))

例

In [267]: random_permutation([1,2,3])
Out[267]: (2, 1, 3)

In [268]: random_permutation([1,2,3])
Out[268]: (3, 2, 1)

・random_combination(iterable,r)
組み合わせをランダムに一つ返します

import random
def random_combination(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    indices = sorted(random.sample(range(n), r))
    return tuple(pool[i] for i in indices)

例

In [280]: random_combination([1,2,3],2)
Out[280]: (1, 3)

In [281]: random_combination([1,2,3],2)
Out[281]: (1, 2)

・random_combination_with_replacement(iterable,r)
重複組み合わせをランダムに一つ返します

def random_combination_with_replacement(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    indices = sorted(random.randrange(n) for i in range(r))
    return tuple(pool[i] for i in indices)

例

In [286]: random_combination([1,2,3],2)
Out[286]: (3, 3)

In [287]: random_combination([1,2,3],2)
Out[287]: (1, 3)