python list comprehensions; compressing a list of lists?

Steve Cooper picture Steve Cooper · Jul 3, 2009 · Viewed 57.4k times · Source

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

Ants Aasma picture Ants Aasma · Jul 3, 2009

You can have nested iterations in a single list comprehension:

[filename for path in dirs for filename in os.listdir(path)]

which is equivalent (at least functionally) to:

filenames = []
for path in dirs:
    for filename in os.listdir(path):
        filenames.append(filename)