Python列表推导创建多个列表

Question

我想创建两个列表listOfA和listOfB存储的索引A和B来自另一个列表秒.

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

输出应该是两个列表

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

我能用两个陈述来做到这一点.

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

但是,我想仅使用列表推导仅在一次迭代中执行此操作.是否可以在一个声明中完成？就像是listOfA,listOfB=[--code goes here--]

Answer 1

列表推导的定义是生成一个列表对象.你的2个列表对象的长度不同; 你必须使用副作用来达到你想要的效果.

不要在这里使用列表推导.只需使用普通循环:

listOfA, listOfB = [], []

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

这使您只需执行一个循环; 这将击败任何两个列表推导,至少直到开发人员找到一种方法使列表推导构建列表的速度是具有单独list.append()调用的循环的两倍.

我想任何一天挑这个在嵌套列表理解仅仅是能够产生在同一行两个列表.正如Python的禅宗所说:

可读性很重要.

Answer 2

有点; 关键是生成一个2元素列表,然后解压缩:

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

也就是说,我认为这样做非常糟糕,显式循环更具可读性.

Answer 3

解决此问题的一个很好的方法是使用defaultdict.正如@Martin所说,列表理解不是生成两个列表的正确工具.使用defaultdict可以使用单次迭代创建隔离.此外,您的代码不会受到任何形式的限制.

>>> from collections import defaultdict
>>> s=['A','B','A','A','A','B','B']
>>> listOf = defaultdict(list)
>>> for idx, elem in enumerate(s):
    listOf[elem].append(idx)
>>> listOf['A'], listOf['B']
([0, 2, 3, 4], [1, 5, 6])