静かなる名辞

pythonとプログラミングのこと

2019/03/22:TechAcademyがteratailの質問・回答を盗用していた件
2019/03/26:TechAcademy盗用事件 公式発表と深まる疑念



【python】itertools.chainを使って複数のiterableを一つにまとめる

概要

 複数のiterable(リストとか)を結合させてループさせたいときがあります。

>>> lst1 = [1, 2, 3]
>>> lst2 = [4, 5, 6]
>>> # 1, 2, 3, 4, 5, 6というループをやりたい

 連結すればできたりしますが、余計なメモリを確保するのでスマートではないし、パフォーマンスが気になります。

>>> for x in lst1 + lst2:
...     print(x)
... 
1
2
3
4
5
6

 というマニアックなお悩みを解決してくれるのがitertools.chainです。

使い方

 リファレンスを見ると使い方が書いてあります。

itertools --- 効率的なループ実行のためのイテレータ生成関数 — Python 3.7.4 ドキュメント

 まあまったく難しいことはなく、

>>> from itertools import chain
>>> for x in chain(lst1, lst2):
...     print(x)
... 
1
2
3
4
5
6

 たったこれだけで使えます。単純ですね。

速度を測る

 実際にパフォーマンスがいいのかどうか試してみましょう。

import timeit
from itertools import chain

def f1(a, b):
    for x in a + b:
        pass

def f2(a, b):
    for x in chain(a, b):
        pass

def main():
    for n in [10, 100, 1000, 10000]:
        print(n)
        a, b = list(range(n)), list(range(n, 2*n))
        t1 = timeit.timeit(lambda : f1(a, b), number=1000)
        t2 = timeit.timeit(lambda : f2(a, b), number=1000)
        print("{0:.6f}\n{1:.6f}".format(t1, t2))
        
if __name__ == "__main__":
    main()

""" =>
10
0.000657
0.000822
100
0.005128
0.004024
1000
0.032347
0.025703
10000
0.249460
0.232873
"""

 なんかどっちでも良いような気が……論理的にはchainを使ったほうがスマートですが、

  • リストの結合と値の取り出しは速い
  • chainはかえって遅い

 というあたりが本質かな、という気がします。

 素のリストとchainの速度も比較してみる。

import timeit
from itertools import chain

def f1(lst):
    for x in lst:
        pass

def f2(lst):
    for x in chain(lst):
        pass

def main():
    lst = list(range(10**4))
    t1 = timeit.timeit(lambda : f1(lst), number=1000)
    t2 = timeit.timeit(lambda : f2(lst), number=1000)
    print("{0:.6f}\n{1:.6f}".format(t1, t2))
    # 0.092083
    # 0.130389

if __name__ == "__main__":
    main()

 かえって遅い可能性もあるということです。

諦めて二重forで書く

 これで良いんじゃ……

import timeit
from itertools import chain

def f1(a, b):
    for x in a + b:
        pass

def f2(a, b):
    for x in chain(a, b):
        pass

def f3(a, b):
    for lst in [a, b]:
        for x in lst:
            pass

def main():
    for n in [10, 100, 1000, 10000]:
        print(n)
        a, b = list(range(n)), list(range(n, 2*n))
        t1 = timeit.timeit(lambda : f1(a, b), number=1000)
        t2 = timeit.timeit(lambda : f2(a, b), number=1000)
        t3 = timeit.timeit(lambda : f3(a, b), number=1000)
        print("{0:.6f}\n{1:.6f}\n{2:.6f}".format(t1, t2, t3))
        
if __name__ == "__main__":
    main()

""" =>
10
0.000675
0.000837
0.000726
100
0.003831
0.003594
0.002471
1000
0.033191
0.025650
0.016758
10000
0.279039
0.225997
0.163289
"""

 こんなのが一番速いです。無理に一つにまとめない方がパフォーマンス上良い、というしょうもない結論に。
 

まとめ

 なんかchainするとパフォーマンス上は微妙な気もしますが、使うと連結よりスマートな感じのコードになるので、好みで使うと良いでしょう。