Pythonでflatten(リストを潰す関数)

Python 30分プログラム

30分プログラム、その595。Pythonでflatten(リストを潰す関数)を作ってみる。
flattenはRubyにあるメノッドでネストされた配列を平坦にしてくれる。

>> [1,2,[3,4]].flatten
=> [1, 2, 3, 4]

これをPythonでやってみよう。せっかくだから、イテレータを持っていれば何でもOKにしたいよね。

使い方

>>> list(flatten([1,2,3]))
[1, 2, 3]

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

# setが
>>> list(flatten([1,2,set([3,4])]))
[1, 2, 3, 4]

ソースコード

ifの条件式を"iter()が成功すること"にすると文字列の処理で死ぬので、"__iter__を持っていること"にした。

#! /usr/bin/python
# -*- mode:python; coding:utf-8 -*-
#
# flatten.py -
#
# Copyright(C) 2009 by mzp
# Author: MIZUNO Hiroki / mzpppp at gmail dot com
# http://howdyworld.org
#
# Timestamp: 2009/06/01 21:05:19
#
# This program is free software; you can redistribute it and/or
# modify it under MIT Lincence.
#

def is_iter(l):
    try:
        iter(l)
        return True
    except TypeError:
        return False

def flatten(xs):
    if '__iter__' in dir(xs):
        for ys in map(flatten,iter(xs)):
            for y in ys:
                yield y
    else:
        yield xs

if __name__ == '__main__':
    print list(flatten([1,2,[3,4]]))
    print list(flatten([1,2,"abc"]))
    print list(flatten([1,2,set([5,5,1])]))

参考