Question

我可以通过将数组与numpy进行比较来创建新数组吗？

我有3个数组（A1, A2, A3）。如何在A1 == 2 and A2 > A3找到所有索引并在新数组中写入值5？

我有这个matlab代码正在执行此操作：

index = find(A1==2 & A2>A3);
new_array(index) = 5;

我找到 putmask 和 logical_and ，但不确定这是否是正确的工具以及如何在我的情况下使用它。谢谢！

Answer 1

以下代码使用&将条件链接在一起。因此，只要A1 == 2和A2 > A3都是True，index数组就会True

import numpy as np

A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3

new_array = np.zeros(len(A1))

index = (A1 == 2) & (A2 > A3)

new_array[index] = 5
# array([ 0.,  0.,  5.,  5.,  0.,  0.])

您当然可以使用np.logical_and。但这限制了您只使用两个条件，而您可以在使用&时有效地链接任意数量。

Answer 2

您可以使用np.where功能

import numpy as np

A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3

out = np.where((A1 == 2) & (A2 >= A3), 5, 0)

或更简单

out = ((A1 == 2) & (A2 >= A3))*5

通过比较数组创建新数组

2 个答案: