python - 如何用条件元素组成列表

Question

我用 python 编程了一段时间，我发现这种语言对程序员非常友好，所以也许有一种我不知道如何用条件元素组成列表的技术。简化的例子是:

# in pseudo code
add_two = True
my_list = [
  "one",
  "two" if add_two,
  "three",
]

基本上，我正在寻找一种方便的方法来创建一个列表，其中包含在某些特定条件下添加的列表。

一些看起来不太好的替代方案:

add_two = True

# option 1
my_list = []
my_list += ["one"]
my_list += ["two"] if add_two else []
my_list += ["three"]


# option 2
my_list = []
my_list += ["one"]
if add_two: my_list += ["two"]
my_list += ["three"]

有什么可以简化的吗？干杯!

Answer 1

如果您可以创建一个 bool 值列表来表示您希望从候选列表中保留哪些元素，则可以非常简洁地完成此操作。例如:

candidates = ['one', 'two', 'three', 'four', 'five']
include = [True, True, False, True, False]
result = [c for c, i in zip(candidates, include) if i]
print(result)
# ['one', 'two', 'four']

如果你可以使用 numpy，这会变得更加简洁:

import numpy as np
candidates = np.array(['one', 'two', 'three', 'four', 'five'])
include = [True, True, False, True, False]
print(candidates[include])  # can use boolean indexing directly!
# ['one', 'two', 'four']

最后，按照评论中的建议，您可以使用 itertools.compress() .请注意，这会返回一个迭代器，因此您必须将其解包。

from itertools import compress
print([v for v in compress(candidates, include)])
# ['one', 'two', 'four']