基于逻辑索引提取子数组

Question

在MATLAB中查看以下代码：

a = [1,2; 5,6]
b = [-1,1; -1,1]
d = a(b(:)>0)

现在d将是2x1数组，[2；6]。这是因为数组b仅在位置(1,2(和(2,2(处具有正条目，并且代码的第三行正在提取这些位置中的a的元素。

Python中有等效的方法可以做到这一点吗？我搜索了numpy文档，但找不到任何文档。在我的实际代码中，我有多个大型多维数组，我希望基于其他数组的元素从中提取元素。当然，这可以通过嵌套for循环来完成，但如果有像MATLAB这样更好的方法，那就更好了。

Answer 1

假设a和b是numpy阵列，则使用：

d = a[b > 0]

在numpy中，索引是用[]运算符完成的。

Answer 2

不使用任何库：

a = [[1, 2], [5, 6]]
b = [[-1, 1], [-1, 1]]
d = [
a_xy
for a_x, b_x in zip(a, b)
for a_xy, b_xy in zip(a_x, b_x)
if b_xy > 0
]

使用numpy:

import numpy as np
a = np.array([[1, 2], [5, 6]])
b = np.array([[-1, 1], [-1, 1]])
d = a[b > 0]