Python

Python 3 で flatten する方法いろいろ

何度も再発明してるのが面倒になったのでまとめておく。

単純な flatten

1階層だけ flatten する方法。

いずれの方法も何らかのモジュールを import する必要があるし、[1, [2, [3]]] みたいな変な形の配列には使えない。

reduce を使うパターン

from functools import reduce

reduce(lambda a, b: a + b, [[1,2,3],[4,5,6],[7,8,9]])  #=> [1, 2, 3, 4, 5, 6, 7, 8, 9]

chain を使うパターン

from itertools import chain

list(chain.from_iterable([[1,2,3],[4,5,6],[7,8,9]]))  #=> [1, 2, 3, 4, 5, 6, 7, 8, 9]

NumPy を使うパターン

import numpy as np

np.array([[1,2,3],[4,5,6],[7,8,9]]).flatten()  #=> array([1, 2, 3, 4, 5, 6, 7, 8, 9])

この場合は当然だけど返ってくるのは NumPy 配列になる。

sum を使うパターン

sum([[1,2,3],[4,5,6],[7,8,9]], [])  #=> [1, 2, 3, 4, 5, 6, 7, 8, 9]

[追記] 当初 sum は「Python 2 では使えたが Python 3 ではできなくなったパターン」と書いていたが、やり方が間違っていただけでそんなことはなかった。何も import しなくてもいいし、1階層だけの flatten ならもうこれでいいのでは。

再帰的 flatten

リスト内包表記を使うと再帰的に flatten できる。

flatten = lambda x: [z for y in x for z in (flatten(y) if hasattr(y, '__iter__') else (y,))]

flatten([[1,2,3],[4,5,6],[7,8,9]])  #=> [1, 2, 3, 4, 5, 6, 7, 8, 9]

この方法は変な形の配列に対しても適用できるので汎用性が高い。

flatten([1,[2,[3]]])  #=> [1, 2, 3]

ただし配列のみならず iterable なオブジェクトをすべて展開してしまうので、配列の中に iterable なオブジェクトを入れている場合は注意する必要がある。