itertools --- 効率的なループ用のイテレータ生成関数群¶
このモジュールは イテレータ を構築する部品を実装しています。プログラム言語 APL, Haskell, SML からアイデアを得ていますが、 Python に適した形に修正されています。
このモジュールは、高速でメモリ効率に優れ、単独でも組合せても使用することのできるツールを標準化したものです。同時に、このツール群は "イテレータの代数" を構成していて、pure Python で簡潔かつ効率的なツールを作れるようにしています。
例えば、SML の作表ツール tabulate(f) は f(0), f(1), ... のシーケンスを作成します。同じことを Python では map() と count() を組合せて map(f, count()) という形で実現できます。
汎用イテレータ:
イテレータ |
引数 |
結果 |
使用例 |
|---|---|---|---|
p [,func] |
p0, p0+p1, p0+p1+p2, ... |
|
|
p, n |
(p0, p1, ..., p_n-1), ... |
|
|
p, q, ... |
p0, p1, ... plast, q0, q1, ... |
|
|
iterable |
p0, p1, ... plast, q0, q1, ... |
|
|
data, selectors |
(d[0] if s[0]), (d[1] if s[1]), ... |
|
|
[start[, step]] |
start, start+step, start+2*step, ... |
|
|
p |
p0, p1, ... plast, p0, p1, ... |
|
|
predicate, seq |
|
|
|
predicate, seq |
predicate(elem) が偽になる |
|
|
iterable[, key] |
key(v) の値でグループ化したサブイテレータ |
|
|
seq, [start,] stop [, step] |
seq[start:stop:step] |
|
|
iterable |
(p[0], p[1]), (p[1], p[2]) |
|
|
elem [,n] |
elem, elem, elem, ... 無限もしくは n 回 |
|
|
func, seq |
func(*seq[0]), func(*seq[1]), ... |
|
|
predicate, seq |
|
|
|
it, n |
it1, it2 , ... itn 一つのイテレータを n 個に分ける |
|
|
p, q, ... |
(p[0], q[0]), (p[1], q[1]), ... |
|
組合せイテレータ:
イテレータ |
引数 |
結果 |
|---|---|---|
p, q, ... [repeat=1] |
デカルト積、ネストしたforループと等価 |
|
p[, r] |
長さrのタプル列、重複なしのあらゆる並び |
|
p, r |
長さrのタプル列、ソートされた順で重複なし |
|
p, r |
長さrのタプル列、ソートされた順で重複あり |
使用例 |
結果 |
|---|---|
|
|
|
|
|
|
|
|
Itertool関数¶
次の関数はいずれもイテレータを構築して返します。 無限長のストリームを提供するものもあるので、その場合はストリームを途中で打ち切る関数やループでのみ使用してください。
- itertools.accumulate(iterable[, function, *, initial=None])¶
累積和、または他の 2 引数関数による累積結果を返すイテレータを作成します。
function のデフォルトは加算です。 function は、累積値と iterable から取り出した値の 2 つの引数を受け取る必要があります。
initial を指定すると、その値から累積を開始し、出力は入力のイテラブルより 1 要素多くなります。
およそ次と等価です:
def accumulate(iterable, function=operator.add, *, initial=None): '累積値を返す' # accumulate([1,2,3,4,5]) → 1 3 6 10 15 # accumulate([1,2,3,4,5], initial=100) → 100 101 103 106 110 115 # accumulate([1,2,3,4,5], operator.mul) → 1 2 6 24 120 iterator = iter(iterable) total = initial if initial is None: try: total = next(iterator) except StopIteration: return yield total for element in iterator: total = function(total, element) yield total
それまでの最小値を求めるには、function に
min()を指定します。 それまでの最大値にはmax()、累積積にはoperator.mul()を指定します。 返済予定表 を作成するには、利息を加算して返済額を差し引きます。>>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8] >>> list(accumulate(data, max)) # それまでの最大値 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9] >>> list(accumulate(data, operator.mul)) # 累積積 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0] # 年利 5% の借入額 1000 に対し、毎年 90 の返済を 10 回行う >>> update = lambda balance, payment: round(balance * 1.05) - payment >>> list(accumulate(repeat(90, 10), update, initial=1_000)) [1000, 960, 918, 874, 828, 779, 728, 674, 618, 559, 497]
最終的な累積値だけを返す類似の関数については
functools.reduce()を見てください。Added in version 3.2.
バージョン 3.3 で変更: 省略可能な function 引数を追加しました。
バージョン 3.8 で変更: オプションの initial パラメータが追加されました。
- itertools.batched(iterable, n, *, strict=False)¶
iterable から得られるデータを n 個ごとに一つのタプルにまとめます。 一番最後のバッチは n 個より少なくなる可能性があります。
strict が真の場合、最後のバッチが n 個未満なら
ValueErrorを送出します。入力のイテラブルを反復処理し、データを最大 n 個ずつタプルにまとめます。 入力は、1 つのバッチを満たすのに必要な分だけ遅延して消費されます。 バッチがいっぱいになるか、入力のイテラブルを 使い切る とすぐに結果を yield します。
>>> flattened_data = ['roses', 'red', 'violets', 'blue', 'sugar', 'sweet'] >>> unflattened = list(batched(flattened_data, 2)) >>> unflattened [('roses', 'red'), ('violets', 'blue'), ('sugar', 'sweet')]
およそ次と等価です:
def batched(iterable, n, *, strict=False): # batched('ABCDEFG', 3) → ABC DEF G if n < 1: raise ValueError('n must be at least one') iterator = iter(iterable) while batch := tuple(islice(iterator, n)): if strict and len(batch) != n: raise ValueError('batched(): incomplete batch') yield batch
Added in version 3.12.
バージョン 3.13 で変更: strict オプションを追加しました。
- itertools.chain(*iterables)¶
最初のイテラブルから 使い切る まで要素を返し、次のイテラブルへ進むことを、すべてのイテラブルを使い切るまで繰り返すイテレータを作成します。 複数のデータ源を単一のイテレータにまとめます。 およそ次のコードと同等です。
def chain(*iterables): # chain('ABC', 'DEF') → A B C D E F for iterable in iterables: yield from iterable
内包表記でのアンパック でも同様のことができ、
list(chain(p, q))は[*s for s in (p, q)]と書けます。
- classmethod chain.from_iterable(iterable)¶
chain()のためのもう一つのコンストラクタです。遅延評価される iterable 引数一つから連鎖した入力を受け取ります。この関数は、以下のコードとほぼ等価です:def from_iterable(iterables): # chain.from_iterable(['ABC', 'DEF']) → A B C D E F for iterable in iterables: yield from iterable
内包表記でのアンパック でも同様のことができ、
list(chain.from_iterable(iterables))は[*s for s in iterables]と書けます。
- itertools.combinations(iterable, r)¶
入力 iterable の要素からなる長さ r の部分列を返します。
出力は
product()の部分列であり、iterable の部分列になる項目だけを残したものです。 出力の長さはmath.comb()で求められ、0 ≤ r ≤ nの場合はn! / r! / (n - r)!、r > nの場合は 0 です。組み合わせのタプルは、入力の iterable の要素の順序に基づく辞書式順序で生成されます。 入力の iterable がソート済みなら、出力のタプルもソートされた順に生成されます。
要素は値ではなく位置によって区別されます。 入力の要素の値が互いに異なれば、各組み合わせの中で値が重複することはありません。
およそ次と等価です:
def combinations(iterable, r): # combinations('ABCD', 2) → AB AC AD BC BD CD # combinations(range(4), 3) → 012 013 023 123 pool = tuple(iterable) n = len(pool) if r > n: return indices = list(range(r)) yield tuple(pool[i] for i in indices) while True: for i in reversed(range(r)): if indices[i] != i + n - r: break else: return indices[i] += 1 for j in range(i+1, r): indices[j] = indices[j-1] + 1 yield tuple(pool[i] for i in indices)
- itertools.combinations_with_replacement(iterable, r)¶
入力 iterable から、それぞれの要素が複数回現れることを許して、長さ r の要素の部分列を返します。
出力は
product()の部分列であり、iterable の部分列 (同じ要素の繰り返しを許す) になる項目だけを残したものです。n > 0の場合、返される部分列の数は(n + r - 1)! / r! / (n - 1)!です。組み合わせのタプルは、入力の iterable の要素の順序に基づく辞書式順序で生成されます。 入力の iterable がソート済みなら、出力のタプルもソートされた順に生成されます。
要素は値ではなく位置によって区別されます。 入力の要素の値が互いに異なれば、生成される組み合わせ同士も重複しません。
およそ次と等価です:
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: return indices = [0] * r yield tuple(pool[i] for i in indices) while True: for i in reversed(range(r)): if indices[i] != n - 1: break else: return indices[i:] = [indices[i] + 1] * (r - i) yield tuple(pool[i] for i in indices)
Added in version 3.1.
- itertools.compress(data, selectors)¶
selectors の対応する要素が真である data の要素を返すイテレータを作成します。 イテラブルの data または selectors のどちらかを 使い切る と終了します。 およそ次のコードと同等です。
def compress(data, selectors): # compress('ABCDEF', [1,0,1,0,1,1]) → A C E F return (datum for datum, selector in zip(data, selectors) if selector)
Added in version 3.1.
- itertools.count(start=0, step=1)¶
start から始まり、一定間隔で値を返すイテレータを作成します。
map()と組み合わせて連続するデータ点を生成したり、zip()と組み合わせて連番を付けたりできます。 およそ次のコードと同等です。def count(start=0, step=1): # count(10) → 10 11 12 13 14 ... # count(2.5, 0.5) → 2.5 3.0 3.5 ... n = start while True: yield n n += step
浮動小数点数で数える場合は、
(start + step * i for i in count())のように乗算を使うコードに置き換えると、精度が向上することがあります。バージョン 3.1 で変更: step 引数が追加され、非整数の引数が許されるようになりました。
- itertools.cycle(iterable)¶
iterable の要素を返すと同時に、それぞれのコピーを保存するイテレータを作成します。 イテラブルを 使い切る と、保存したコピーから要素を返します。 これを無限に繰り返します。 およそ次のコードと同等です。
def cycle(iterable): # cycle('ABCD') → A B C D A B C D A B C D ... saved = [] for element in iterable: yield element saved.append(element) while saved: for element in saved: yield element
このイテレータツールは、イテラブルの長さによっては大量の補助記憶領域を必要とする場合があります。
- itertools.dropwhile(predicate, iterable)¶
predicate が真を返す間は iterable の要素を捨て、その後はすべての要素を返すイテレータを作成します。 およそ次のコードと同等です。
def dropwhile(predicate, iterable): # dropwhile(lambda x: x<5, [1,4,6,3,8]) → 6 3 8 iterator = iter(iterable) for x in iterator: if not predicate(x): yield x break for x in iterator: yield x
predicateが初めて偽を返すまでは 一切 出力しないため、このイテレータツールは最初の出力までに時間がかかる場合があります。
- itertools.filterfalse(predicate, iterable)¶
iterable の要素を絞り込み、predicate が偽を返す要素だけを返すイテレータを作成します。 predicate が
Noneの場合、要素自体が偽のものを返します。 およそ次のコードと同等です。def filterfalse(predicate, iterable): # filterfalse(lambda x: x<5, [1,4,6,3,8]) → 6 8 if predicate is None: predicate = bool for x in iterable: if not predicate(x): yield x
- itertools.groupby(iterable, key=None)¶
iterable 内で同じキーを持つ連続した要素をグループにまとめ、キーとグループを返すイテレータを作成します。 key は各要素に対するキー値を計算する関数です。 キーを指定しない場合や
Noneにした場合、key 関数のデフォルトは恒等関数になり要素をそのまま返します。 通常、iterable は同じキー関数でソート済みである必要があります。groupby()の操作は Unix のuniqフィルターと似ています。key関数の値が変わるたびに区切りを入れ、新しいグループを生成します (このために通常同じkey関数でソートしておく必要があるのです)。 この動作は SQL の入力順に関係なく共通の要素を集約する GROUP BY とは違います。返されるグループはそれ自体がイテレータで、
groupby()と iterable を共有しています。もととなる iterable を共有しているため、groupby()オブジェクトの要素取り出しを先に進めると、それ以前の要素であるグループは見えなくなってしまいます。従って、データが後で必要な場合にはリストの形で保存しておく必要があります:groups = [] uniquekeys = [] data = sorted(data, key=keyfunc) for k, g in groupby(data, keyfunc): groups.append(list(g)) # グループのイテレータをリストとして保存する uniquekeys.append(k)
groupby()はおよそ次と等価です:def groupby(iterable, key=None): # [k for k, g in groupby('AAAABBBCCDAABBB')] → A B C D A B # [list(g) for k, g in groupby('AAAABBBCCD')] → AAAA BBB CC D keyfunc = (lambda x: x) if key is None else key iterator = iter(iterable) exhausted = False def _grouper(target_key): nonlocal curr_value, curr_key, exhausted yield curr_value for curr_value in iterator: curr_key = keyfunc(curr_value) if curr_key != target_key: return yield curr_value exhausted = True try: curr_value = next(iterator) except StopIteration: return curr_key = keyfunc(curr_value) while not exhausted: target_key = curr_key curr_group = _grouper(target_key) yield curr_key, curr_group if curr_key == target_key: for _ in curr_group: pass
- itertools.islice(iterable, stop)¶
- itertools.islice(iterable, start, stop[, step])
iterableから選択した要素を返すイテレータを作成します。 シーケンスのスライスのように動作しますが、start、stop、step に負の値は指定できません。start が 0 または
Noneの場合、位置 0 から反復を開始します。 それ以外の場合は、start に達するまでiterableの要素を読み飛ばします。stop が
Noneの場合、入力を 使い切る まで反復を続けます (入力が尽きなければ終了しません)。 それ以外の場合は、指定した位置で終了します。step が
Noneの場合、間隔のデフォルト値は 1 です。 step を 1 より大きくして要素を読み飛ばすようにしなければ、要素を連続して返します。およそ次と等価です:
def islice(iterable, *args): # islice('ABCDEFG', 2) → A B # islice('ABCDEFG', 2, 4) → C D # islice('ABCDEFG', 2, None) → C D E F G # islice('ABCDEFG', 0, None, 2) → A C E G s = slice(*args) start = 0 if s.start is None else s.start stop = s.stop step = 1 if s.step is None else s.step if start < 0 or (stop is not None and stop < 0) or step <= 0: raise ValueError indices = count() if stop is None else range(max(start, stop)) next_i = start for i, element in zip(indices, iterable): if i == next_i: yield element next_i += step
If the input is an iterator, then fully consuming the islice advances the input iterator by
max(start, stop)steps regardless of the step value.
- itertools.pairwise(iterable)¶
入力の iterable から隣り合う要素のペアを順に返します。 連続するペアは要素を 1 つ共有します。
出力イテレータに含まれる 2 要素タプルの数は、入力の要素数より 1 つ少なくなります。 入力のイテラブルの要素が 2 つ未満なら、出力は空になります。
およそ次と等価です:
def pairwise(iterable): # pairwise('ABCDEFG') → AB BC CD DE EF FG iterator = iter(iterable) a = next(iterator, None) for b in iterator: yield a, b a = b
Added in version 3.10.
- itertools.permutations(iterable, r=None)¶
iterable の要素からなる長さ r の 順列 を順に返します。
r が指定されない場合や
Noneの場合、r はデフォルトで iterable の長さとなり、可能な最長の順列の全てが生成されます。出力は
product()の部分列であり、同じ要素を繰り返し使う項目を取り除いたものです。 出力の長さはmath.perm()で求められ、0 ≤ r ≤ nの場合はn! / (n - r)!、r > nの場合は 0 です。順列のタプルは、入力の iterable の要素の順序に基づく辞書式順序で生成されます。 入力の iterable がソート済みなら、出力のタプルもソートされた順に生成されます。
要素は値ではなく位置によって区別されます。 入力の要素の値が互いに異なれば、各順列の中で値が重複することはありません。
およそ次と等価です:
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)) cycles = list(range(n, n-r, -1)) yield tuple(pool[i] for i in indices[:r]) while n: for i in reversed(range(r)): cycles[i] -= 1 if cycles[i] == 0: indices[i:] = indices[i+1:] + indices[i:i+1] cycles[i] = n - i else: j = cycles[i] indices[i], indices[-j] = indices[-j], indices[i] yield tuple(pool[i] for i in indices[:r]) break else: return
- itertools.product(*iterables, repeat=1)¶
入力のイテラブルの デカルト積 です。
ジェネレータ式の入れ子になった for ループとおよそ等価です。たとえば
product(A, B)は((x,y) for x in A for y in B)と同じものを返します。入れ子ループは走行距離計と同じように右端の要素がイテレーションごとに更新されていきます。このパターンは辞書式順序を作り出し、入力のイテレート可能オブジェクトたちがソートされていれば、直積タプルもソートされた順に出てきます。
イテラブル自身との直積を計算するためには、オプションの repeat キーワード引数に繰り返し回数を指定します。たとえば
product(A, repeat=4)はproduct(A, A, A, A)と同じ意味です。この関数は以下のコードとおよそ等価ですが、実際の実装ではメモリ中に中間結果を作りません:
def product(*iterables, 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 if repeat < 0: raise ValueError('repeat argument cannot be negative') pools = [tuple(pool) for pool in iterables] * repeat result = [[]] for pool in pools: result = [x+[y] for x in result for y in pool] for prod in result: yield tuple(prod)
product()は動作する前に、入力のイテラブルを完全に読み取り、直積を生成するためにメモリ内に値を蓄えます。したがって、入力が有限の場合に限り有用です。
- itertools.repeat(object[, times])¶
object を繰り返し返すイテレータを作成します。 times 引数を指定しなければ、無限に繰り返します。
およそ次と等価です:
def repeat(object, times=None): # repeat(10, 3) → 10 10 10 if times is None: while True: yield object else: for i in range(times): yield object
repeat は map や zip に定数のストリームを与えるためによく利用されます:
>>> list(map(pow, range(10), repeat(2))) [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
- itertools.starmap(function, iterable)¶
iterable から取り出した引数を使って function を計算するイテレータを作成します。 引数がすでにタプルにまとめられている場合に、
map()の代わりに使用します。map()とstarmap()の違いは、function(a,b)とfunction(*c)の違いに対応します。 およそ次のコードと同等です。def starmap(function, iterable): # starmap(pow, [(2,5), (3,2), (10,3)]) → 32 9 1000 for args in iterable: yield function(*args)
- itertools.takewhile(predicate, iterable)¶
predicate が真を返す間、iterable の要素を返すイテレータを作成します。 およそ次のコードと同等です。
def takewhile(predicate, iterable): # takewhile(lambda x: x<5, [1,4,6,3,8]) → 1 4 for x in iterable: if not predicate(x): break yield x
条件を最初に満たさなくなった要素も入力イテレータから消費され、後から取り出すことはできません。 そのため、takewhile を 使い切った 後で、入力イテレータの残りをさらに消費したいアプリケーションでは問題になることがあります。 この問題を回避するには、代わりに more-itertools の before_and_after() を使うことを検討してください。
- itertools.tee(iterable, n=2)¶
一つの iterable から n 個の独立したイテレータを返します。
およそ次と等価です:
def tee(iterable, n=2): if n < 0: raise ValueError if n == 0: return () iterator = _tee(iterable) result = [iterator] for _ in range(n - 1): result.append(_tee(iterator)) return tuple(result) class _tee: def __init__(self, iterable): it = iter(iterable) if isinstance(it, _tee): self.iterator = it.iterator self.link = it.link else: self.iterator = it self.link = [None, None] def __iter__(self): return self def __next__(self): link = self.link if link[1] is None: link[0] = next(self.iterator) link[1] = [None, None] value, self.link = link return value
入力の iterable がすでに tee イテレータオブジェクトの場合、返されるタプルの全要素は、その入力を生成した元の
tee()呼び出しで作られたかのように構築されます。 この「平坦化」により、入れ子になったtee()呼び出しは下層の同じデータの連鎖を共有し、呼び出しを連鎖させずに 1 回の更新で済むようになります。この平坦化の性質によって、tee イテレータの次の要素を効率よく先読みできます。
def lookahead(tee_iterator): "入力を進めずに次の値を返す" [forked_iterator] = tee(tee_iterator, 1) return next(forked_iterator)
>>> iterator = iter('abcdef') >>> [iterator] = tee(iterator, 1) # 入力を先読みできるようにする >>> next(iterator) # イテレータを進める 'a' >>> lookahead(iterator) # 次の値を確認する 'b' >>> next(iterator) # 引き続き先に進める 'b'
teeイテレータはスレッドセーフではありません。 元の iterable がスレッドセーフであっても、同じtee()呼び出しで返されたイテレータを同時に使用すると、RuntimeErrorが送出される場合があります。このイテレータツールは、一時的に保存する必要があるデータ量によっては、多くの補助記憶領域を必要とすることがあります。 一般に、一方のイテレータが他方の処理開始前にデータの大半またはすべてを使う場合は、
tee()よりlist()を使う方が高速です。
- itertools.zip_longest(*iterables, fillvalue=None)¶
iterables の各イテラブルから取り出した要素をまとめるイテレータを作成します。
イテラブルの長さが異なる場合、不足する値を fillvalue で埋めます。 fillvalue のデフォルトは
Noneです。最も長いイテラブルを 使い切る まで反復を続けます。
およそ次と等価です:
def zip_longest(*iterables, fillvalue=None): # zip_longest('ABCD', 'xy', fillvalue='-') → Ax By C- D- iterators = list(map(iter, iterables)) num_active = len(iterators) if not num_active: return while True: values = [] for i, iterator in enumerate(iterators): try: value = next(iterator) except StopIteration: num_active -= 1 if not num_active: return iterators[i] = repeat(fillvalue) value = fillvalue values.append(value) yield tuple(values)
入力のイテラブルのいずれかが無限長の可能性がある場合は、
islice()やtakewhile()などでzip_longest()をラップし、呼び出し回数を制限してください。
Itertools レシピ¶
この節では、既存の itertools を素材としてツールセットを拡張するためのレシピを示します。
itertools レシピの主な目的は学習です。
レシピは、個々のツールを捉えるさまざまな考え方を示します。
例えば、chain.from_iterable は平坦化という概念に関係します。
また、starmap() と repeat() の組み合わせ方など、ツールを組み合わせるためのアイデアも示します。
さらに、itertools を operator や collections モジュールと組み合わせる使い方や、map()、filter()、reversed()、enumerate() などの組み込みの反復処理ツールと組み合わせる使い方も示します。
レシピのもう 1 つの目的は、新しい機能を育てることです。
accumulate()、compress()、pairwise() は、当初はレシピとして登場しました。
現在は sliding_window()、derangements()、sieve() のレシピで、その有用性を検証しています。
これらのレシピのほぼすべてと、さらに多くのレシピを、Python Package Index の more-itertools プロジェクトからインストールできます。
python -m pip install more-itertools
多くのレシピは、基盤となるツール群と同じ高い性能を提供します。 イテラブル全体を一度にメモリに読み込まず、要素を 1 つずつ処理することで、優れたメモリ効率を保ちます。 ツールを 関数型のスタイル でつなぐことで、コードの量も少なくできます。 インタープリタのオーバーヘッドが生じる for ループや ジェネレータ よりも、「ベクトル化された」構成要素を優先して使うことで、高速性を維持します。
from itertools import (accumulate, batched, chain, combinations, compress,
count, cycle, filterfalse, groupby, islice, permutations, product,
repeat, starmap, tee, zip_longest)
from collections import Counter, deque
from contextlib import suppress
from functools import reduce
from heapq import heappush, heappushpop, heappush_max, heappushpop_max
from math import comb, isqrt, prod, sumprod
from operator import getitem, is_not, itemgetter, mul, neg, truediv
# ==== 基本的な 1 行の処理 ====
def take(n, iterable):
"iterable の先頭 n 個の要素をリストとして返す。"
return list(islice(iterable, n))
def prepend(value, iterable):
"iterable の先頭に値を 1 つ追加する。"
# prepend(1, [2, 3, 4]) → 1 2 3 4
return chain([value], iterable)
def repeatfunc(function, times=None, *args):
"指定した引数で function の呼び出しを繰り返す。"
if times is None:
return starmap(function, repeat(args))
return starmap(function, repeat(args, times))
def flatten(list_of_lists):
"入れ子を 1 段階平坦化する。"
return chain.from_iterable(list_of_lists)
def ncycles(iterable, n):
"シーケンスの要素を n 回繰り返して返す。"
return chain.from_iterable(repeat(tuple(iterable), n))
def loops(n):
"""n 回ループする。
range(n) と似ているが整数を作らない。
"""
# for _ in loops(100): ...
return repeat(None, n)
def tail(n, iterable):
"最後の n 個の要素を反復するイテレータを返す。"
# tail(3, 'ABCDEFG') → E F G
return iter(deque(iterable, maxlen=n))
def consume(iterator, n=None):
"""iterator を n 個先に進める。
n が None の場合はすべて消費する。
"""
# イテレータを C の速度で消費する関数を使う。
if n is None:
deque(iterator, maxlen=0)
else:
next(islice(iterator, n, n), None)
def nth(iterable, n, default=None):
"n 番目の要素、またはデフォルト値を返す。"
return next(islice(iterable, n, None), default)
def quantify(iterable, predicate=bool):
"True または False を返す predicate を適用し、True の数を数える。"
return sum(map(predicate, iterable))
def first_true(iterable, default=False, predicate=None):
"最初の真の値を返し、真の値がなければ default を返す。"
# first_true([a, b, c], x) → a or b or c or x
# first_true([a, b], x, f) → a if f(a) else b if f(b) else x
return next(filter(predicate, iterable), default)
def all_equal(iterable, key=None):
"すべての要素が互いに等しければ True を返す。"
# all_equal('4٤௪౪໔', key=int) → True
return len(take(2, groupby(iterable, key))) <= 1
# ==== データパイプライン ====
def unique_justseen(iterable, key=None):
"""順序を保ち、連続する重複を除いた要素を yield する。
直前の要素だけを記憶する。
"""
# unique_justseen('AAAABBBCCDAABBB') → A B C D A B
# unique_justseen('ABBcCAD', str.casefold) → A B c A D
if key is None:
return map(itemgetter(0), groupby(iterable))
return map(next, map(itemgetter(1), groupby(iterable, key)))
def unique_everseen(iterable, key=None):
"""順序を保ち、重複を除いた要素を yield する。
それまでのすべての要素を記憶する。
"""
# unique_everseen('AAAABBBCCDAABBB') → A B C D
# unique_everseen('ABBcCAD', str.casefold) → A B c D
seen = set()
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
def unique(iterable, key=None, reverse=False):
"""重複を除いた要素をソート順で yield する。
ハッシュ不可能な入力もサポートする。
"""
# unique([[1, 2], [3, 4], [1, 2]]) → [1, 2] [3, 4]
sequenced = sorted(iterable, key=key, reverse=reverse)
return unique_justseen(sequenced, key=key)
def sliding_window(iterable, n):
"データを互いに重なり合う固定長のまとまりに集める。"
# sliding_window('ABCDEFG', 3) → ABC BCD CDE DEF EFG
iterator = iter(iterable)
window = deque(islice(iterator, n - 1), maxlen=n)
for x in iterator:
window.append(x)
yield tuple(window)
def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
"データを互いに重ならない固定長のまとまりに集める。"
# grouper('ABCDEFG', 3, fillvalue='x') → ABC DEF Gxx
# grouper('ABCDEFG', 3, incomplete='strict') → ABC DEF ValueError
# grouper('ABCDEFG', 3, incomplete='ignore') → ABC DEF
iterators = [iter(iterable)] * n
match incomplete:
case 'fill':
return zip_longest(*iterators, fillvalue=fillvalue)
case 'strict':
return zip(*iterators, strict=True)
case 'ignore':
return zip(*iterators)
case _:
raise ValueError('Expected fill, strict, or ignore')
def roundrobin(*iterables):
"すべての入力イテラブルを使い切るまで、順番に巡回する。"
# roundrobin('ABC', 'D', 'EF') → A D E B F C
# アルゴリズムの考案者: George Sakkis
iterators = map(iter, iterables)
for num_active in range(len(iterables), 0, -1):
iterators = cycle(islice(iterators, num_active))
yield from map(next, iterators)
def subslices(seq):
"シーケンスの空でない連続した部分スライスをすべて返す。"
# subslices('ABCD') → A AB ABC ABCD B BC BCD C CD D
slices = starmap(slice, combinations(range(len(seq) + 1), 2))
return map(getitem, repeat(seq), slices)
def derangements(iterable, r=None):
"不動点を持たない長さ r の順列を生成する。"
# derangements('ABCD') → BADC BCDA BDAC CADB CDAB CDBA DABC DCAB DCBA
# アルゴリズムの考案者: Stefan Pochmann
seq = tuple(iterable)
pos = tuple(range(len(seq)))
have_moved = map(map, repeat(is_not), repeat(pos), permutations(pos, r=r))
valid_derangements = map(all, have_moved)
return compress(permutations(seq, r=r), valid_derangements)
def iter_index(iterable, value, start=0, stop=None):
"シーケンスまたはイテラブル内で value が現れるインデックスを返す。"
# iter_index('AABCADEAF', 'A') → 0 1 4 7
seq_index = getattr(iterable, 'index', None)
if seq_index is None:
iterator = islice(iterable, start, stop)
for i, element in enumerate(iterator, start):
if element is value or element == value:
yield i
else:
stop = len(iterable) if stop is None else stop
i = start
with suppress(ValueError):
while True:
yield (i := seq_index(value, i, stop))
i += 1
def iter_except(function, exception, first=None):
"例外が起きるまで呼び出すインターフェースをイテレータに変換する。"
# iter_except(d.popitem, KeyError) → 非ブロッキングの辞書イテレータ
with suppress(exception):
if first is not None:
yield first()
while True:
yield function()
# ==== 数学的演算 ====
def multinomial(*counts):
"多重集合の相異なる並べ方の数。"
# Counter('abracadabra').values() → 5 2 2 1 1
# multinomial(5, 2, 2, 1, 1) → 83160
return prod(map(comb, accumulate(counts), counts))
def powerset(iterable):
"iterable の部分列を短い順に返す。"
# powerset([1,2,3]) → () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)
s = list(iterable)
return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
def sum_of_squares(iterable):
"入力値の二乗の和を求める。"
# sum_of_squares([10, 20, 30]) → 1400
return sumprod(*tee(iterable))
# ==== 行列の操作 ====
def reshape(matrix, columns):
"2 次元の行列を指定した列数に変形する。"
# reshape([(0, 1), (2, 3), (4, 5)], 3) → (0, 1, 2) (3, 4, 5)
return batched(chain.from_iterable(matrix), columns, strict=True)
def transpose(matrix):
"2 次元の行列の行と列を入れ替える。"
# transpose([(1, 2, 3), (11, 22, 33)]) → (1, 11) (2, 22) (3, 33)
return zip(*matrix, strict=True)
def matmul(m1, m2):
"2 つの行列の積を求める。"
# matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]) → (49, 80) (41, 60)
n = len(m2[0])
return batched(starmap(sumprod, product(m1, transpose(m2))), n)
# ==== 多項式の演算 ====
def convolve(signal, kernel):
"""2 つのイテラブルの離散線形畳み込み。
多項式の乗算に相当する。
畳み込みは数学的には交換可能だが、入力の評価方法は異なる。
signal は遅延して消費され、無限長でもよい。
kernel は計算が始まる前にすべて消費される。
記事: https://betterexplained.com/articles/intuitive-convolution/
動画: https://www.youtube.com/watch?v=KuXjwB4LzSA
"""
# convolve([1, -1, -20], [1, -3]) → 1 -4 -17 60
# convolve(data, [0.25, 0.25, 0.25, 0.25]) → 移動平均 (ぼかし)
# convolve(data, [1/2, 0, -1/2]) → 1 階微分の推定
# convolve(data, [1, -2, 1]) → 2 階微分の推定
kernel = tuple(kernel)[::-1]
n = len(kernel)
padded_signal = chain(repeat(0, n-1), signal, repeat(0, n-1))
windowed_signal = sliding_window(padded_signal, n)
return map(sumprod, repeat(kernel), windowed_signal)
def polynomial_from_roots(roots):
"""多項式の根から係数を求める。
(x - 5) (x + 4) (x - 3) を展開すると: x³ -4x² -17x + 60
"""
# polynomial_from_roots([5, -4, 3]) → [1, -4, -17, 60]
factors = zip(repeat(1), map(neg, roots))
return list(reduce(convolve, factors, [1]))
def polynomial_eval(coefficients, x):
"""指定した値における多項式の値を求める。
ホーナー法より数値的に安定した計算を行う。
"""
# x = 5 での x³ -4x² -17x + 60 の値を求める
# polynomial_eval([1, -4, -17, 60], x=5) → 0
n = len(coefficients)
if not n:
return type(x)(0)
powers = map(pow, repeat(x), reversed(range(n)))
return sumprod(coefficients, powers)
def polynomial_derivative(coefficients):
"""多項式の 1 階導関数を求める。
f(x) = x³ -4x² -17x + 60
f'(x) = 3x² -8x -17
"""
# polynomial_derivative([1, -4, -17, 60]) → [3, -8, -17]
n = len(coefficients)
powers = reversed(range(1, n))
return list(map(mul, coefficients, powers))
# ==== 整数論 ====
def sieve(n):
"n 未満の素数。"
# sieve(30) → 2 3 5 7 11 13 17 19 23 29
if n > 2:
yield 2
data = bytearray((0, 1)) * (n // 2)
for p in iter_index(data, 1, start=3, stop=isqrt(n) + 1):
data[p*p : n : p+p] = bytes(len(range(p*p, n, p+p)))
yield from iter_index(data, 1, start=3)
def factor(n):
"n の素因数。"
# factor(99) → 3 3 11
# factor(1_000_000_000_000_007) → 47 59 360620266859
# factor(1_000_000_000_000_403) → 1000000000000403
for prime in sieve(isqrt(n) + 1):
while not n % prime:
yield prime
n //= prime
if n == 1:
return
if n > 1:
yield n
def is_prime(n):
"n が素数なら True を返す。"
# is_prime(1_000_000_000_000_403) → True
return n > 1 and next(factor(n)) == n
def totient(n):
"n 以下の自然数のうち、n と互いに素なものの数。"
# https://mathworld.wolfram.com/TotientFunction.html
# totient(12) → 4 (len([1, 5, 7, 11]) == 4 のため)
for prime in set(factor(n)):
n -= n // prime
return n
# ==== 累積統計量 ====
def running_mean(iterable):
"それまでの値の平均。"
# running_mean([37, 33, 38, 28]) → 37 35 36 34
return map(truediv, accumulate(iterable), count(1))
def running_min(iterable):
"それまでの値の最小値。"
# running_min([37, 33, 38, 28]) → 37 33 33 28
return accumulate(iterable, func=min)
def running_max(iterable):
"それまでの値の最大値。"
# running_max([37, 33, 38, 28]) → 37 37 38 38
return accumulate(iterable, func=max)
def running_median(iterable):
"それまでの値の中央値。"
# running_median([37, 33, 38, 28]) → 37 35 37 35
read = iter(iterable).__next__
lo = [] # 最大ヒープ
hi = [] # 要素数が lo と同じか 1 つ少ない最小ヒープ
with suppress(StopIteration):
while True:
heappush_max(lo, heappushpop(hi, read()))
yield lo[0]
heappush(hi, heappushpop_max(lo, read()))
yield (lo[0] + hi[0]) / 2
def running_statistics(iterable):
"それまでの値について統計量をまとめる。"
# タプルを生成する: (個数, 最小値, 中央値, 最大値, 平均)
t0, t1, t2, t3 = tee(iterable, 4)
return zip(
count(1),
running_min(t0),
running_median(t1),
running_max(t2),
running_mean(t3),
)