Edited at

pythonでflatten

Rubyにはネストされた配列を1次元配列にフラット化するflattenというメソッドがあるのですが、pythonでも使いたくなることがあるので、どんな方法があるか調べてみました。


やりたいこと

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

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


Qiitaで見つけた方法

http://qiita.com/kento1218@github/items/f3baf574aadb3d1cbeae


ジェネレータ関数にする

Qiitaで見つけた方法を参考にしてジェネレータ関数にしてみました。

リスト化するには list(flatten([[1, 2], [3, [4, 5]]])) のようにする必要があります。

def flatten(data):

for item in data:
if hasattr(item, '__iter__'):
for element in flatten(item):
yield element
else:
yield item


ジェネレータ内包表記にする

def flatten(data):

return (element
for item in data
for element in (flatten(item) if hasattr(item, '__iter__') else [item]))


リスト内包表記にする

def flatten(data):

return [element
for item in data
for element in (flatten(item) if hasattr(item, '__iter__') else [item])]


既存ライブラリにある関数を使う

ここで見つけました。

http://stackoverflow.com/questions/2158395/flatten-an-irregular-list-of-lists-in-python

from compiler.ast import flatten


sum関数を使ってリストのリストを深さ1だけflattenする

>>> data = [[1, 2], [3], [4, 5, [6]]]

>>> sum(data, [])
[1, 2, 3, 4, 5, [6]]


深さ指定オプション付き関数

Rubyではフラット化する深さを指定できるので、対応してみました。


flatten.py

#!/usr/bin/env python

# -*- coding:utf-8 -*-

def flatten(data, depth=-1):
"""
flatten(data) -> list
flatten(data, depth) -> list

Return flatted data of list or tupple as list.

>>> data = [[1, 2], [3, [4, 5, [6]]]]
>>> flatten(data)
[1, 2, 3, 4, 5, 6]
>>> flatten(data, 0)
[[1, 2], [3, [4, 5, [6]]]]
>>> flatten(data, 1)
[1, 2, 3, [4, 5, [6]]]
>>> flatten(data, 2)
[1, 2, 3, 4, 5, [6]]
>>> flatten(data, 3)
[1, 2, 3, 4, 5, 6]
"""
return [element
for item in data
for element in (flatten(item, depth - 1)
if depth != 0 and hasattr(item, '__iter__')
else [item])
]


$ python -m doctest -v flatten.py

Trying:
data = [[1, 2], [3, [4, 5, [6]]]]
Expecting nothing
ok
Trying:
flatten(data)
Expecting:
[1, 2, 3, 4, 5, 6]
ok
Trying:
flatten(data,0)
Expecting:
[[1, 2], [3, [4, 5, [6]]]]
ok
Trying:
flatten(data, 1)
Expecting:
[1, 2, 3, [4, 5, [6]]]
ok
Trying:
flatten(data, 2)
Expecting:
[1, 2, 3, 4, 5, [6]]
ok
Trying:
flatten(data, 3)
Expecting:
[1, 2, 3, 4, 5, 6]
ok
1 items had no tests:
flatten
1 items passed all tests:
6 tests in flatten.flatten
6 tests in 2 items.
6 passed and 0 failed.
Test passed.