通過比較數組創建新陣列

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