python list comprehensions; compressing a list of

Question

guys. I'm trying to find the most elegant solution to a problem and wondered if python has anything built-in for what I'm trying to do.

What I'm doing is this. I have a list, A, and I have a function f which takes an item and returns a list. I can use a list comprehension to convert everything in A like so;

[f(a) for a in A]

But this return a list of lists;

[a1,a2,a3] => [[b11,b12],[b21,b22],[b31,b32]]

What I really want is to get the flattened list;

[b11,b12,b21,b22,b31,b32]

Now, other languages have it; it's traditionally called flatmap in functional programming languages, and .Net calls it SelectMany. Does python have anything similar? Is there a neat way to map a function over a list and flatten the result?

The actual problem I'm trying to solve is this; starting with a list of directories, find all the subdirectories. so;

import os
dirs = ["c:\\usr", "c:\\temp"]
subs = [os.listdir(d) for d in dirs]
print subs

currentliy gives me a list-of-lists, but I really want a list.

Answer 1

The question proposed flatmap. Some implementations are proposed but they may unnecessary creating intermediate lists. Here is one implementation that's base on iterators.

def flatmap(func, *iterable):
    return itertools.chain.from_iterable(map(func, *iterable))

In [148]: list(flatmap(os.listdir, ['c:/mfg','c:/Intel']))
Out[148]: ['SPEC.pdf', 'W7ADD64EN006.cdr', 'W7ADD64EN006.pdf', 'ExtremeGraphics', 'Logs']

In Python 2.x, use itertools.map in place of map.

Answer 2

subs = []
map(subs.extend, (os.listdir(d) for d in dirs))

(but Ants's answer is better; +1 for him)

Answer 3

You could try itertools.chain(), like this:

import itertools
import os
dirs = ["c:\\usr", "c:\\temp"]
subs = list(itertools.chain(*[os.listdir(d) for d in dirs]))
print subs

itertools.chain() returns an iterator, hence the passing to list().

Answer 4

Google brought me next solution:

def flatten(l):
   if isinstance(l,list):
      return sum(map(flatten,l))
   else:
      return l

Answer 5

def flat_list(arr):
    send_back = []
    for i in arr:
        if type(i) == list:
            send_back += flat_list(i)
        else:
            send_back.append(i)
    return send_back

Answer 6

You can use pyxtension:

from pyxtension.streams import stream
stream([ [1,2,3], [4,5], [], [6] ]).flatMap() == range(7)