逐元素布尔乘法

Question

我正在寻找最简单美观的方法来对这些列表进行乘法运算：

a = [True, False, True, False, False]

b = [100, 200]

a 中 True 元素的数量等于 b 的长度 我需要的答案是这里的 [100, 0, 200, 0, 0]。

有没有简单的方法得到这个答案？看起来像是逐元素相乘，但实际上第二个列表的大小较小，所以常规的numpy方法不起作用，需要进行复杂的编码。

希望您能找到一个好的解决方案。

- John Makkonen

---

你所需要的答案是一个“列表”，但你已经标记了“numpy”，它通常处理“numpy.ndarray”对象，这是两种具有不同性能权衡/特征的数据结构。那么，到底是哪个？ - juanpa.arrivillaga

Answer 1

你可以使用numpy来完成此操作：

c = np.array(a).astype(int)

c[c==1] = b

>>> c
array([100,   0, 200,   0,   0])

注意 如果您需要按照所需输出将结果作为列表而不是numpy数组返回，请使用 c.tolist()

- sacuL

Answer 2

有一种基于迭代器的相当简洁的解决方案

it = iter(b)
[next(it) if x else 0 for x in a]
# [100, 0, 200, 0, 0]

- Patrick Haugh

---

可能是最好的非numpy解决方案。 - AGN Gazer

可能是最好的解决方案。 - Alexander

@Alexander 不，这并不是最好的“终点”。当数据已经是数组时，基于numpy的解决方案更快。 - AGN Gazer

Answer 3

另一种选择是使用 numpy.place() 或 numpy.put()：

c = np.zeros_like(a, dtype=np.result_type(b))
np.place(c, a, b)

或者

c = np.zeros_like(a, dtype=np.result_type(b))
np.put(c, np.flatnonzero(a), b)

我更喜欢第一个选项。

时间：

生成更大的数据集：

In [39]: a = np.random.randint(0, 2, 100000).astype(np.bool_)

In [40]: b = np.random.randint(0, 100000, np.sum(a))

@sacul的解决方案：

In [41]: %timeit c = np.array(a).astype(int); c[c==1] = b
621 µs ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

我的第一个选项：

In [42]: %timeit c = np.zeros_like(a, dtype=np.result_type(b)); np.place(c, a, b)
462 µs ± 8.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

- AGN Gazer

Answer 4

从上面的@AGN中借鉴一些内容，涉及编程，保留HTML，不解释。

c = np.zeros_like(a, dtype = b.dtype)
c[a] = b

- Daniel F