Python list comprehensions to create multiple lists [duplicate]

Question

I want to create two lists listOfA and listOfB to store indices of A and B from another list s.

s=['A','B','A','A','A','B','B'] 

Output should be two lists

listOfA=[0,2,3,4] listOfB=[1,5,6] 

I am able to do this with two statements.

listOfA=[idx for idx,x in enumerate(s) if x=='A'] listOfB=[idx for idx,x in enumerate(s) if x=='B'] 

However, I want to do it in only one iteration using list comprehensions only. Is it possible to do it in a single statement? something like listOfA,listOfB=[--code goes here--]

Answer 1

The very definition of a list comprehension is to produce one list object. Your 2 list objects are of different lengths even; you'd have to use side-effects to achieve what you want.

Don't use list comprehensions here. Just use an ordinary loop:

listOfA, listOfB = [], []  for idx, x in enumerate(s):     target = listOfA if x == 'A' else listOfB     target.append(idx) 

This leaves you with just one loop to execute; this will beat any two list comprehensions, at least not until the developers find a way to make list comprehensions build a list twice as fast as a loop with separate list.append() calls.

I'd pick this any day over a nested list comprehension just to be able to produce two lists on one line. As the Zen of Python states:

Readability counts.

Answer 2

Sort of; the key is to generate a 2-element list that you can then unpack:

listOfA, listOfB = [[idx for idx, x in enumerate(s) if x == c] for c in 'AB'] 

That said, I think it's pretty daft to do it that way, an explicit loop is much more readable.