RubyのflattenをPythonでやってみた

訂正：再帰の上限を増やせることが判明

SICP的な再帰

Ruby - Wikipediaの例をPythonでやってみた−前半戦 - 牌語備忘録 - pygoで書いたコードと同じ

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

def flatten(seq):
    if not seq:
        return []
    if isinstance(seq[0], list):
        return flatten(seq[0]) + flatten(seq[1:])
    return [seq[0]] + flatten(seq[1:])
print list(flatten(a))

実行結果

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

リストの要素を増やして再度実行してみた

a = [1,[2,[3,[4,[5,6],7],8],9],10] * 1000

実行結果

RuntimeError: maximum recursion depth exceeded

エラーになった(´･ω･`)
Pythonの再帰には上限があるらしい。

再帰の上限を増やす

setrecursionlimitを使うと再帰の上限を増やせる。（morchinさんのコメントから。ありがとうございます(・∀・)）

import sys
sys.setrecursionlimit(100000)

上記を追加して再度実行

[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5,
:
大杉なので略
:
...6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

いけた(・∀・)b

yieldをつかってみる

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

def flatten_yield(seq):
    for elem in seq:
        if isinstance(elem, list):
            [(yield s) for s in flatten_yield(elem)]
        else:
            yield elem
print list(flatten_yield(a))

実行結果

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

またリストの要素を増やして再度実行してみた

a = [1,[2,[3,[4,[5,6],7],8],9],10] * 1000

実行結果

[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5,
:
大杉なので略
:
...6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

問題なく処理できた(・∀・)

いちおう速度を計ってみた

import timeit
t= timeit.Timer(setup="a = [1,[2,[3,[4,[5,6],7],8],9],10]*100",
                stmt="""
def flatten(seq):
    if not seq:
        return []
    if isinstance(seq[0], list):
        return flatten(seq[0]) + flatten(seq[1:])
    return [seq[0]] + flatten(seq[1:])
list(flatten(a))
""")
print "A recursive :", t.timeit(number = 100)
t= timeit.Timer(setup="a = [1,[2,[3,[4,[5,6],7],8],9],10]*100",
                stmt="""
def flatten_yield(seq):
    for elem in seq:
        if isinstance(elem, list):
            [(yield s) for s in flatten_yield(elem)]
        else:
            yield elem
list(flatten_yield(a))
""")
print "B yield     :", t.timeit(number = 100)

実行結果

A recursive : 0.740335941315
B yield     : 0.328200817108

yieldを使った方が速いらしい。