Python 3 の些細な不満の一つは flatten が欲しいときに毎回自分で再発明しないといけないところ
— ほと (@hoto17296) June 13, 2018
何度も再発明してるのが面倒になったのでまとめておく。
まとめ
1階層だけならこれでいい (ついでに型もつけた)
def flatten[T](arr: list[list[T]]) -> list[T]:
return sum(arr, [])
単純な flatten
１階層だけ 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 しなくてもいいし、１階層だけの flatten ならもうこれでいいのでは。
再帰的 flatten
リスト内包表記を使うと再帰的に flatten できる。
flatten = lambda x: [z for y in x for z in (flatten(y) if hasattr(y, '__iter__') and not isinstance(y, str) else (y,))]
flatten([[1,2,3],[4,5,6],[7,8,9]]) #=> [1, 2, 3, 4, 5, 6, 7, 8, 9]
この方法は変な形の配列に対しても適用できるので汎用性が高い。
[追記] 「文字列が含まれると無限再帰になってしまう」と指摘いただいたので and not isinstance(y, str) を追加した。
flatten([1,[2,[3]]]) #=> [1, 2, 3]
ただし配列のみならず iterable なオブジェクトをすべて展開してしまうので、配列の中に iterable なオブジェクトを入れている場合は注意する必要がある。