numpy多个布尔索引数组

Question

我有一个要在其上使用布尔索引的数组，其中包含多个索引数组，每个索引数组生成一个不同的数组。示例：

w = np.array([1,2,3])
b = np.array([[False, True, True], [True, False, False]])

应该按照以下方式返回：

[[2,3], [1]]

我假设由于包含True的单元格的数量在蒙版之间可以变化，所以我不能期望结果驻留在2d numpy数组中，但是我仍然希望有比遍历蒙版更优雅的东西。通过第i个w掩码为其索引b的结果。

我是否缺少更好的选择？

编辑：我接下来要做的下一步是对w[b]返回的每个数组求和，返回一个标量列表。如果这样可以使问题更容易解决，我也想知道。

Answer 1

假设您想要一个numpy数组列表，您可以简单地使用一个理解：

w = np.array([1,2,3])
b = np.array([[False, True, True], [True, False, False]])
[w[bool] for bool in b]

# [array([2, 3]), array([1])]

如果您的目标只是所使用的掩盖值的总和：

np.sum(w*b) # 6

或

np.sum(w*b, axis=1) # array([5, 1])
# or b @ w

…因为False次，您的数字将为0，因此不会影响总和。

Answer 2

尝试一下：

[w[x] for x in b]

希望这会有所帮助。