创建多个列表的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 = [[idx for idx, x in enumerate(s) if x == c] for c in 'AB']

也就是说，我认为这样做是相当愚蠢的，一个显式循环更具可读性。

Answer 2

列表理解的定义就是生成一个列表对象。你的两个列表对象的长度甚至不同；你必须使用副作用来实现你想要的。

这里不要使用列表理解。使用普通循环：

listOfA, listOfB = [], []

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

这只剩下一个循环来执行；这将击败任何两个列表理解，至少在开发人员找到方法使列表理解构建一个列表的速度比使用单独的list.append()调用的循环快两倍之前是这样。

任何时候，我都会选择一个嵌套的列表理解只需就可以在一行上生成两个列表。如Zen of Python所述：

Readability counts.

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