views:

1134

answers:

8

Hi, 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.

+3  A: 

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().

Steef
+5  A: 
subs = []
map(subs.extend, (os.listdir(d) for d in dirs))

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

RichieHindle
Using reduce (or sum, which saves you many characters and an import;-) for this is just wrong -- you keep uselessly tossing away old lists to make a new one for each d. @Ants has the right answer (smart of @Steve to accept it!).
Alex Martelli
@Alex: Good point. Fixed.
RichieHindle
+10  A: 

You can find a good answer in itertools' recipes:

def flatten(listOfLists):
    return list(chain.from_iterable(listOfLists))

(Note: requires Python 2.6+)

Roberto Liffredo
+3  A: 

Google brought me next solution:

def flatten(l):
   if isinstance(l,list):
      return sum(map(flatten,l))
   else:
      return l
Vestel
Would be a little better if it handled generator expressions too, and would be a lot better if you explained how to use it...
ephemient
+17  A: 

You can have nested iterations in a single list comprehension:

[filename for path in dirs for filename in os.listdir(path)]
Ants Aasma
+8  A: 

You could just do the straightforward:

subs = []
for d in dirs:
    subs.extend(os.listdir(d))
Anon
Yep, this is fine (though not quite as good as @Ants') so I'm giving it a +1 to honor its simplicity!
Alex Martelli
+2  A: 

You can concatenate lists using the normal addition operator:

>>> [1, 2] + [3, 4]
[1, 2, 3, 4]

The built-in function sum will add the numbers in a sequence and can optionally start from a specific value:

>>> sum(xrange(10), 100)
145

Combine the above to flatten a list of lists:

>>> sum([[1, 2], [3, 4]], [])
[1, 2, 3, 4]

You can now define your flatmap:

>>> def flatmap(f, seq):
...   return sum([f(s) for s in seq], [])
... 
>>> flatmap(range, [1,2,3])
[0, 0, 1, 0, 1, 2]

Edit: I just saw the critique in the comments for another answer and I guess it is correct that Python will needlessly build and garbage collect lots of smaller lists with this solution. So the best thing that can be said about it is that it is very simple and concise if you're used to functional programming :-)

Martin Geisler
+1  A: 
>>> listOfLists = [[1, 2],[3, 4, 5], [6]]
>>> reduce(list.__add__, listOfLists)
[1, 2, 3, 4, 5, 6]

I'm guessing the itertools solution is more efficient than this, but this feel very pythonic and avoids having to import a library just for the sake of a single list operation.

Julian