路傍のプログラマ

只のプログラマが綴る愚痴と備忘録

flattening list in Python

2009-01-30 20:16:26 | プログラミング
入れ子になったリストを平らにならす、という問題。
必要になったのでやってみます。

この問題、どうもかなり古典的な問題らしく、いろんな人が解いてます。

私が必要としたのは、とにかく深い入れ子を処理できること。

具体的な方針としては、
・スタックを使わない。スタックの制限はsys.setrecursionlimit()で緩和できるけど、インタプリタのグローバルな状態に影響を与えるのはあんまりうれしくない。
・2乗のアルゴリズムを使わない。リストの中程に要素を挿入したりとか、pop(0)とか、そういったたぐいの操作を行わない。

参考にしたのは、
http://markmail.org/message/tbz3s6a3wj42cfxs
とか
http://rightfootin.blogspot.com/2006/09/more-on-python-flatten.html
です。

できたコードは、これ(もしコピーして使われる方がいらっしゃるのなら、全角スペースを半角スペースになおしてください)。

def flatten(l, ltypes=(list, tuple)):
  ltype = type(l)
  r = []
  stk = [iter(l)]
  while stk:
    curIter = stk[-1]
    for item in curIter:
      if isinstance(item, ltypes):
        stk.append(iter(item))
        break # for item
      r.append(item)
    else:
      stk.pop()
  return ltype(r)

深く入れ子になったリストを使ってベンチマークしましたが、調子いいですよ。