Pythonでリストをflattenする方法まとめ

参考

1段ネストしたリストをflattenする方法

2段以上ネストしたリストをflattenする方法

色々方法があるみたいなので書留め。

1段ネストしたリストをflattenするとは

つまり、

l = [[1,2,3],[4,5,6], [7], [8,9]]

[1,2,3,4,5,6,7,8,9]

にするということ。

reduceを使った方法

reduceを使ってリストを畳み込むことができる。

reduce(lambda x,y: x+y,l)

組込み関数を使ったほうが読みやすいでしょう。

from operator import add
reduce(add, l)

sumによるやり方

空のリストに要素を積んでいく。これはちょっと思いつかなかった。

sum(l, [])

ここまではreduceの動作に慣れ親しんでいれば容易に理解できる。

リスト内包表記

sumやreduceによる方法は遅いらしい。参考リンク先の説明によると中間結果を毎回すべてコピーしているからだそうだ。
リスト内包表記は無駄なリストを作らないので高速に動作する。

[item for sublist in l for item in sublist]

ただしこれは入れ子の内包表記なのでかなり読みづらい。
かといって2段ネストさせても面白くないし…

result = []
for sublist in l:
    for item in sublist:
        result.append(item)

extend

extendのほうが内包表記より速かったりする。

x = []
for s in l:
    x.extend(s)

extendメソッドをキャッシュしておくとさらに速くなる。

ex = [].extend
for s in l:
    ex(s)

itertools.chain

実はpython.orgにはitertoolsのレシピとしてflattenする方法が載っていたりする。

9.7. itertools ― 効率的なループ実行のためのイテレータ生成関数 ― Python 2.7ja1 documentation

しかもこれが一番速い。

from itertools import chain
list(chain.from_iterable(l))

StackOverFlowばかり見てないで本家ドキュメントを読みなさいというお告げが聞こえる。

1段ネストしたリストのflattenまとめ

ここまで上で挙げた方法について、

[[1,2,3],[4,5,6], [7], [8,9]]**99

をTimer.timeit(1000)でflattenするのにかかった時間を計測し、グラフ化したものを載せる。右側ほど速い。

2段以上ネストしたリストのflatten

これまでのflattenでは多重ネストしたリストの場合内側のリストが残ってしまうという問題(というか仕様)があった。

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

また、

[1, [2, 3]]

のようなリストを平坦化しようとすると、1がIterableでないため、

TypeError: 'int' object is not iterable

となってしまう。これらの不都合な動作を抑制して、

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

と動作するflatten関数がほしい。

再帰による実装

任意の深さのリストをトラバースするとなると再帰による方法が自然に思える。
しかしflatten呼び出しの条件はどのように考えたら良いだろうか。
抽象基底クラスを用いて変数がIterableであるかどうかを判断し、その結果によって処理を振り分けるのが良いだろう。

isinstance(element, collections.Iterable)

抽象基底クラスが何であるか、なぜcollection.Iterableなのかは以下を参照。

用語集 ― Python 2.7ja1 documentation

ところで、文字列もIterableであるがこれは再帰されては困るので弾いておこう。

isinstance(element, collections.Iterable) and not isinstance(element, basestring)

これで再帰する条件は分かったのでflattenは実装できるだろう。

def flatten(nested_list):
    result = []
    for element in nested_list:
        if isinstance(element, collections.Iterable) and not isinstance(element, basestring):
            result.extend(flatten(element))
        else:
            result.append(element)
    return result

ジェネレータを使えばより効率良く実行できるだろうか。

def flatten_gen(l):
    for el in l:
        if isinstance(el, collections.Iterable) and not isinstance(el, basestring):
            for sub in flatten_gen(el):
                yield sub
        else:
            yield el

再帰なしの実装

さて、単純な再帰を使った実装では呼び出しの深さに限界が存在する。

l = []
for i in range(1000):
    l = [l, i]
flatten(l)
# => Runtime Error: maximum recursion depth exceeded

たいていの場合、再帰による実装はループとスタックを使った実装に置き換えることができる。

def flatten_loop(l):
    i = 0
    while i < len(l):
        while isinstance(l[i], collections.Iterable):
            if not l[i]:
                l.pop(i)
                i -= 1
                break
            else:
                l[i:i + 1] = l[i]
        i += 1
    return l

こいつはどんなに深くネストしていてもリストのサイズが許す限り平坦化を行うことができる。

2段以上ネストしたリストのflattenまとめ

l = []
for i in range(100):
    l = [l, i]

について再帰、ジェネレータ再帰、ループの3つの方法によってflattenした結果を示す(Timer.timeit(1000))。
ジェネレータの方が遅いという結果。条件が悪い?