Displaying posts written in

3月 2010

Python  | 2010 年 3 月 31 日 By | No comments yet

Pythonでリストのflatten・高階関数を使って

Python には組み込みの flatten がない。そこで次のように実装した。 

def traverse(f, x):
    if isinstance(x, (list, tuple)):
        for e in x:
            traverse(f, e)
    else:
        f(x)
 
def flatten(listOfList):
    arr = []
    traverse(arr.append, listOfList)
    return arr

「Python flatten」でググると様々な実装例が見つかるのだが、妙に凝った実装が多くて、こういうシンプルな実装はあまり見当たらなかった。

ただしこの実装は再帰を使っているので、ネストが深すぎるとオーバーフローする。 

x = []
for i in xrange(1000):
    x = [x, i]
print flatten(x)
# => RuntimeError: maximum recursion depth exceeded

オーバーフローを避けるには traverse をループを使って書き下す。Pythonでは末尾再帰の最適化は行われない。 

def traverse(f, x, isSeq = lambda y: isinstance(y, (list, tuple))):
    if not isSeq(x): return f(x)
    pos = depth = 0
    stack = { depth: [x, pos, len(x)] }
    while depth >= 0:
        lis, pos, size = stack[depth]
        if pos < size:
            stack[depth][1] = pos + 1
            if isSeq(lis[pos]):
                depth += 1
                stack[depth] = [lis[pos], 0, len(lis[pos])]
            else:
                f(lis[pos])
        else:
            depth -= 1

これで100000とかネストしても平気。 

x = []
for i in xrange(100000):
    x = [x, i]
print flatten(x)
# =>
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
...省略...
, 99977, 99978, 99979, 99980, 99981, 99982, 99983, 99984, 99985, 99986, 99987, 99988, 
99989, 99990, 99991, 99992, 99993, 99994, 99995, 99996, 99997, 99998, 99999]