Skip to content
Advertisement

How can I get a flat result from a list comprehension instead of a nested list?

I have a list A, and a function f which takes an item of A and returns a list. I can use a list comprehension to convert everything in A like [f(a) for a in A], but this returns a list of lists. Suppose my input is [a1,a2,a3], resulting in [[b11,b12],[b21,b22],[b31,b32]].

How can I get the flattened list [b11,b12,b21,b22,b31,b32] instead? In other words, in Python, how can I get what is traditionally called flatmap in functional programming languages, or SelectMany in .NET?

(In the actual code, A is a list of directories, and f is os.listdir. I want to build a flat list of subdirectories.)


See also: How do I make a flat list out of a list of lists? for the more general problem of flattening a list of lists after it’s been created.

Advertisement

Answer

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)
Advertisement