python 多次元配列を1次元配列に変換する

前回まで分割キーボードを作っていたが完成しやっと手になじんできた。

2次元以上の配列を1次元のリストへ

# 配列返すバージョン
def flatten(x, notype=(str, bytes, bytearray), _recur=False):
    if _recur is False and isinstance(x, notype) or not hasattr(x, "__iter__"):
        return [x]
    return [z for y in x for z in (flatten(y, _recur=True)
              if y is not None and hasattr(y, "__iter__") and not isinstance(y, notype) else (y,))]

# generatorバージョン
def iterflatten(x, notype=(str, bytes, bytearray), _recur=False):
    if _recur is False and isinstance(x, notype) or not hasattr(x, "__iter__"):
        yield x
    yield from (z for y in x for z in (flatten(y, _recur=True)
              if y is not None and hasattr(y, "__iter__") and not isinstance(y, notype) else (y,)))

どんなにネストが深かろうと、1次元配列を返す。
追記 more_itertoolsにcollapseという関数があることを知った。

from more_itertools import collapse
>>> list(collapse([1,2,[3,4,[5,6,[7]]]]))
[1, 2, 3, 4, 5, 6, 7]

>>> flatten([1,2,[3,4,[5,6,[7]]]])
[1, 2, 3, 4, 5, 6, 7]

>>> flatten([1,2])
[1, 2]

# iterableならなんでも回して返す
>>> flatten([1,2,[range(3)]])
[1, 2, 0, 1, 2]

>>> flatten("1,2")
["1,2"]

>>> flatten(None)
[None]

>>> flatten([])
[]

_2d_10000 = [[0, 1, 2]] * 10000

>>> %%timeit
... sum(_2d_10000, [])
466 ms ± 81.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

import itertools

>>> %%timeit
... list(itertools.chain.from_iterable(_2d_10000))
356 µs ± 5.44 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

>>> %%timeit
... flatten(_2d_10000)
9.79 ms ± 31.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 2次元配列ならitertoolsのfrom_iterableより遅いが、sumより速く、3次元配列以上に対応できる汎用性優位



# 追記 3次元以上
>>> _3nest_10000 = [1,2,[3,4,[5,6,[7]]]] * 10000
>>> from more_itertools import collapse
>>> %%timeit
list(collapse(_3nest_10000))
74.5 ms ± 233 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> %%timeit
flatten(_3nest_10000)
27.7 ms ± 29.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)