itertools – 効率的なループ処理のためのイテレータ関数

目的:itertools モジュールは繰り返し可能な(シーケンスのような)データセットを扱う関数を提供する
利用できるバージョン:2.3

“lazy functional programming language” の Haskell や SML の類似機能からアイディアを得た関数が提供されています。それらの関数は高速且つ省メモリですが、イテレーションベースのアルゴリズムを複雑化させてしまうこともあります。

イテレータベースのコードの方が複数の理由でリストを使用するコードよりも推奨されます。イテレータで扱うデータは必要になるまで生成されないので、全てのデータを同時にメモリ内には格納されません。メモリ使用量が減少するということはパフォーマンスに影響を与えるスワッピングやその他の巨大なデータセットを扱うときの副作用を小さくします。

イテレータのマージと分割

chain() 関数は引数として複数のイテレータを受け取って、それらが1つのシーケンスであるかのように、全てのコンテンツを生成する1つのイテレータを返します。

from itertools import *

for i in chain([1, 2, 3], ['a', 'b', 'c']):
    print i
$ python itertools_chain.py
1
2
3
a
b
c

izip() 関数は引数で受け取った複数のイテレータの要素をタプルに結合するイテレータを返します。それはリストの代わりにイテレータを返すということを除けばビルトイン関数 zip() のように動作します。

from itertools import *

for i in izip([1, 2, 3], ['a', 'b', 'c']):
    print i
$ python itertools_izip.py
(1, 'a')
(2, 'b')
(3, 'c')

islice() 関数は入力イテレータからインデックスで選択された要素を返すイテレータを返します。それはリストのスライシングのように開始、停止、ステップといった同じ引数を取ります。開始とステップの引数はオプションです。

from itertools import *

print 'Stop at 5:'
for i in islice(count(), 5):
    print i

print 'Start at 5, Stop at 10:'
for i in islice(count(), 5, 10):
    print i

print 'By tens to 100:'
for i in islice(count(), 0, 100, 10):
    print i
$ python itertools_islice.py
Stop at 5:
0
1
2
3
4
Start at 5, Stop at 10:
5
6
7
8
9
By tens to 100:
0
10
20
30
40
50
60
70
80
90

tee() 関数はオリジナルの入力イテレータを1つ受け取って、複数の独立したイテレータを返します(デフォルトは2つ)。それはその入力イテレータから読み込まれた値を繰り返してファイルや標準出力に書き込む Unix の tee ユーティリティによく似た動作をします。

from itertools import *

r = islice(count(), 5)
i1, i2 = tee(r)

for i in i1:
    print 'i1:', i
for i in i2:
    print 'i2:', i
$ python itertools_tee.py
i1: 0
i1: 1
i1: 2
i1: 3
i1: 4
i2: 0
i2: 1
i2: 2
i2: 3
i2: 4

tee() が生成する新たなイテレータはそのオリジナルの入力イテレータを共有するので、生成後はオリジナルの入力イテレータを使用してはいけません。もしオリジナルの入力イテレータから値を取り出すと、新たなイテレータはその取り出された値を生成しません。

from itertools import *

r = islice(count(), 5)
i1, i2 = tee(r)

for i in r:
    print 'r:', i
    if i > 1:
        break

for i in i1:
    print 'i1:', i
for i in i2:
    print 'i2:', i
$ python itertools_tee_error.py
r: 0
r: 1
r: 2
i1: 3
i1: 4
i2: 3
i2: 4

入力を変換する

imap() 関数は入力イテレータの値を引数にして関数を呼び出すイテレータを返します。そして呼び出されたその関数の結果を返します。それは入力イテレータの値がなくなったときに停止するいうこと(代わりに None を追加する)を除けば map() 関数のように動作します。

最初のサンプルでは lambda 関数は入力値を2倍します。2番目のサンプルでは lambda 関数は引数で渡される2つのイテレータから受け取った2つの引数の乗算を行います。そしてオリジナルの引数と算出された値のタプルを返します。

from itertools import *

print 'Doubles:'
for i in imap(lambda x:2*x, xrange(5)):
    print i

print 'Multiples:'
for i in imap(lambda x,y:(x, y, x*y), xrange(5), xrange(5,10)):
    print '%d * %d = %d' % i
$ python itertools_imap.py
Doubles:
0
2
4
6
8
Multiples:
0 * 5 = 0
1 * 6 = 6
2 * 7 = 14
3 * 8 = 24
4 * 9 = 36

starmap() 関数は imap() 関数によく似ていますが、複数のイテレータを持つタプルを構築する代わりにマッピングする関数への引数として * 構文で1つのイテレータ内にある要素を分割して渡します。imap() にマッピングされる関数が f(i1, i2) のように引数を渡して呼び出す場合 starmap() にマッピングされる関数は f(*i) になります。

from itertools import *

values = [(0, 5), (1, 6), (2, 7), (3, 8), (4, 9)]
for i in starmap(lambda x,y:(x, y, x*y), values):
    print '%d * %d = %d' % i
$ python itertools_starmap.py
0 * 5 = 0
1 * 6 = 6
2 * 7 = 14
3 * 8 = 24
4 * 9 = 36

新たな値を生成する

count() 関数は無限に連続した整数を生成するイテレータを返します。開始の数のデフォルトはゼロで引数として渡すことができます。上限の引数はありません(結果セットの詳細な制御はビルトイン関数 xrange() を参照)。このサンプルでは、リストの引数の要素がなくなるとそのイテレーションが停止します。

from itertools import *

for i in izip(count(1), ['a', 'b', 'c']):
    print i
$ python itertools_count.py
(1, 'a')
(2, 'b')
(3, 'c')

cycle() 関数は引数のコンテンツを無限に繰り返すイテレータを返します。入力イテレータの完全なコンテンツを覚えておく必要があるので、そのイテレータが大きい場合に少しメモリを消費する可能性があります。このサンプルでは数回サイクルした後でループを脱出するためにカウンタ変数が使用されます。

from itertools import *

i = 0
for item in cycle(['a', 'b', 'c']):
    i += 1
    if i == 10:
        break
    print (i, item)
$ python itertools_cycle.py
(1, 'a')
(2, 'b')
(3, 'c')
(4, 'a')
(5, 'b')
(6, 'c')
(7, 'a')
(8, 'b')
(9, 'c')

repeat() 関数はアクセスされると毎回同じ値を生成するイテレータを返します。オプションの引数で上限値を渡さない限り無限に値を返します。

from itertools import *

for i in repeat('over-and-over', 5):
    print i
$ python itertools_repeat.py
over-and-over
over-and-over
over-and-over
over-and-over
over-and-over

定数が他のイテレータからの値と一緒に含まれる必要があるとき izip() 又は imap()repeat() を結合すると便利です。

from itertools import *

for i, s in izip(count(), repeat('over-and-over', 5)):
    print i, s
$ python itertools_repeat_izip.py
0 over-and-over
1 over-and-over
2 over-and-over
3 over-and-over
4 over-and-over
from itertools import *

for i in imap(lambda x,y:(x, y, x*y), repeat(2), xrange(5)):
    print '%d * %d = %d' % i
$ python itertools_repeat_imap.py
2 * 0 = 0
2 * 1 = 2
2 * 2 = 4
2 * 3 = 6
2 * 4 = 8

フィルタリング

dropwhile() 関数は条件が最初に False になった後で入力イテレータの要素を返すイテレータを返します。その条件が False になった後では各要素がフィルタされず、入力イテレータの残りの要素が返されます。

from itertools import *

def should_drop(x):
    print 'Testing:', x
    return (x<1)

for i in dropwhile(should_drop, [ -1, 0, 1, 2, 3, 4, 1, -2 ]):
    print 'Yielding:', i
$ python itertools_dropwhile.py
Testing: -1
Testing: 0
Testing: 1
Yielding: 1
Yielding: 2
Yielding: 3
Yielding: 4
Yielding: 1
Yielding: -2

dropwhile() とは逆の機能で takewhile() 関数はテスト関数が True を返す限り入力イテレータから要素を返すイテレータを返します。

from itertools import *

def should_take(x):
    print 'Testing:', x
    return (x<2)

for i in takewhile(should_take, [ -1, 0, 1, 2, 3, 4, 1, -2 ]):
    print 'Yielding:', i
$ python itertools_takewhile.py
Testing: -1
Yielding: -1
Testing: 0
Yielding: 0
Testing: 1
Yielding: 1
Testing: 2

ifilter() 関数は、テスト関数が True を返すときのみ要素を含める、ビルトイン関数の filter() のように動作するイテレータを返します。 dropwhile() との違いは全要素に対してテスト関数で判定を行ってその要素を返します。

from itertools import *

def check_item(x):
    print 'Testing:', x
    return (x<1)

for i in ifilter(check_item, [ -1, 0, 1, 2, 3, 4, 1, -2 ]):
    print 'Yielding:', i
$ python itertools_ifilter.py
Testing: -1
Yielding: -1
Testing: 0
Yielding: 0
Testing: 1
Testing: 2
Testing: 3
Testing: 4
Testing: 1
Testing: -2
Yielding: -2

ifilter() とは逆の機能で ifilterfalse() 関数はテスト関数が False を返すときのみ要素を含めるイテレータを返します。

from itertools import *

def check_item(x):
    print 'Testing:', x
    return (x<1)

for i in ifilterfalse(check_item, [ -1, 0, 1, 2, 3, 4, 1, -2 ]):
    print 'Yielding:', i
$ python itertools_ifilterfalse.py
Testing: -1
Testing: 0
Testing: 1
Yielding: 1
Testing: 2
Yielding: 2
Testing: 3
Yielding: 3
Testing: 4
Yielding: 4
Testing: 1
Yielding: 1
Testing: -2

データのグループ化

groupby() 関数は共通のキーでグループ化された値セットを生成するイテレータを返します。

標準ライブラリドキュメントで紹介されているこのサンプルは同じ値を持つ辞書のキーをグループ化します。

from itertools import *
from operator import itemgetter

d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
di = sorted(d.iteritems(), key=itemgetter(1))
for k, g in groupby(di, key=itemgetter(1)):
    print k, map(itemgetter(0), g)
$ python itertools_groupby.py
1 ['a', 'c', 'e']
2 ['b', 'd', 'f']
3 ['g']

これは複数の属性をベースに関連する値をグループ化することを紹介した複雑なサンプルです。期待した通りに動作するグルーピングのために入力シーケンスがキーの順番通りにソートされている必要があることに注意してください。

from itertools import *

class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y
    def __repr__(self):
        return 'Point(%s, %s)' % (self.x, self.y)
    def __cmp__(self, other):
        return cmp((self.x, self.y), (other.x, other.y))

# Point インスタンスのデータセットを作成する
data = list(imap(Point, 
                 cycle(islice(count(), 3)), 
                 islice(count(), 10),
                 )
            )
print 'Data:', data
print

# X 値ベースのソートされていないデータをグループ化しようとする
print 'Grouped, unsorted:'
for k, g in groupby(data, lambda o:o.x):
    print k, list(g)
print

# データをソートする
data.sort()
print 'Sorted:', data
print

# X 値ベースのソートされたデータをグループ化する
print 'Grouped, sorted:'
for k, g in groupby(data, lambda o:o.x):
    print k, list(g)
print
$ python itertools_groupby_seq.py
Data: [Point(0, 0), Point(1, 1), Point(2, 2), Point(0, 3), Point(1, 4), Point(2, 5), Point(0, 6), Point(1, 7), Point(2, 8), Point(0, 9)]

Grouped, unsorted:
0 [Point(0, 0)]
1 [Point(1, 1)]
2 [Point(2, 2)]
0 [Point(0, 3)]
1 [Point(1, 4)]
2 [Point(2, 5)]
0 [Point(0, 6)]
1 [Point(1, 7)]
2 [Point(2, 8)]
0 [Point(0, 9)]

Sorted: [Point(0, 0), Point(0, 3), Point(0, 6), Point(0, 9), Point(1, 1), Point(1, 4), Point(1, 7), Point(2, 2), Point(2, 5), Point(2, 8)]

Grouped, sorted:
0 [Point(0, 0), Point(0, 3), Point(0, 6), Point(0, 9)]
1 [Point(1, 1), Point(1, 4), Point(1, 7)]
2 [Point(2, 2), Point(2, 5), Point(2, 8)]

See also

itertools
本モジュールの標準ライブラリドキュメント
The Standard ML Basis Library
SML のライブラリ
Definition of Haskell and the Standard Libraries
関数型言語 Haskell の標準ライブラリ仕様
Bookmark and Share