け日記

最近はPythonでいろいろやってます

Python: more-itertoolsで複雑なイテレーションを簡単に実装する

開発や分析で「タプルのリストからタプルの1つ目の要素だけを取り出したい」「リストの要素を3つずつ処理したい」といったことがちょくちょく起こります。Pythonでこうしたケースに便利なライブラリ mote_itertools を紹介します。

github.com

インストール

今回は7.0.0を使います。

$ pip install -U more-itertools
Successfully installed more-itertools-7.0.0

実装

今回紹介するメソッドをインポートしておきます。APIの詳細はこちらです。

from more_itertools import chunked, flatten, distribute, divide, split_at, split_before, split_after, unzip, windowed

リストからx個ずつ取り出したい (chunked)

chunked_iter = chunked([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 4)
print(list(chunked_iter))  # [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]]

リストをx個へ均等に分割したい (divide, distributed)

iter1, iter2, iter3 = list(divide(3, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
print(list(iter1), list(iter2), list(iter3))  # [0, 1, 2, 3] [4, 5, 6] [7, 8, 9]

x個置きに分割する場合は、distributedを使うと良いです。

iter1, iter2, iter3 = list(distribute(3, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
print(list(iter1), list(iter2), list(iter3))  # [0, 3, 6, 9] [1, 4, 7] [2, 5, 8]

リストのリスト・タプルのリストを1つのリストに展開したい (flatten)

flatten_iter = flatten([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]])  # [(0, 1, 2, 3), (4, 5, 6, 7), (8, 9)] でもOK
print(list(flatten_iter)) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

リスト内の特定の要素で分割したい (split_at, split_before, split_after)

split_atを使い、4で割り切れる要素で分割しています。分割に使った要素 (0, 4, 8) は除去されます。

splitted_iters = split_at([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], lambda e: e%4 == 0)
print(list(splitted_iters))  # [[], [1, 2, 3], [5, 6, 7], [9]]

分割に使った要素を残す場合、split_beforeまたはsplit_afterを使います。分割した要素が、split_beforeでは先頭、split_afterでは末尾になるように分割されます。

splitted_iters = split_before([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], lambda e: e%4 == 0)
print(list(splitted_iters))  # [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]]

splitted_iters = split_after([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], lambda e: e%4 == 0)
print(list(splitted_iters))  # [[0], [1, 2, 3, 4], [5, 6, 7, 8], [9]]

タプルのリストからタプル要素ごとのリストにしたい (unzip)

要素が無い場合、以降の要素が取得できなくなりますので、注意です。

name_iter, age_iter = unzip([('tanaka', 28), ('suzuki', 33), ('yamada', 23), ('sato', )])
print(list(name_iter))  # ['tanaka', 'suzuki', 'yamada', 'sato']
print(list(age_iter))  # [28, 33, 23]

リストから範囲をずらしながら取り出したい (windowed)

4個ずつずらしながらイテレーションしてます。

windowed_iter = windowed([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 4)
print(list(windowed_iter))  # [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5), (3, 4, 5, 6), (4, 5, 6, 7), (5, 6, 7, 8), (6, 7, 8, 9)]

ずらし幅はstepで指定できます。

windowed_iter = windowed([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 4, step=2)
print(list(windowed_iter))  # [(0, 1, 2, 3), (2, 3, 4, 5), (4, 5, 6, 7), (6, 7, 8, 9)]

まとめ

同じことは二重forループやlambdaを組み合わせれば実現できるのですが、コードの見通しが悪くなってしまってなんだかなあ、という問題を解決できるので、おすすめのライブラリです。