当条件为2时,有效地更新numpy数组

Question

假设您有一个包含3或4个元素的numpy数组

例如:

3
4
4
4
3
3
3

我想改变这个数组,这样如果一个元素是3,那么它应该变成一个数字X,如果元素是4,它应该变成一个数字Y.保证X不同于Y.对于上面的数组我们会得到:

X
Y
Y
Y
X
X
X

我在考虑做这样的事情:

arr[arr==3]=X
arr[arr==4]=Y

但是,如果X是4呢？然后最后整个数组将只包含Ys.

我试图避免出于性能原因使用for循环,但如果这是唯一的方法,我可以负担得起遵循该路线.

Answer 1

由于输入1D数组仅包含3或4,您可以使用np.where,如此 -

np.where(A==3,X,Y)

样品运行#1 -

In [55]: A
Out[55]: array([4, 3, 4, 3, 3, 3, 4, 4, 4, 4, 3, 3, 4, 3, 4])

In [56]: X = 30; Y = 40

In [57]: np.where(A==3,X,Y)
Out[57]: array([40, 30, 40, 30, 30, 30, 40, 40, 40, 40, 30, 30, 40, 30, 40])

样品运行#2(如果X是4) -

In [60]: A
Out[60]: array([4, 3, 4, 3, 3, 3, 4, 4, 4, 4, 3, 3, 4, 3, 4])

In [61]: X = 4; Y = 40

In [62]: np.where(A==3,X,Y)
Out[62]: array([40,  4, 40,  4,  4,  4, 40, 40, 40, 40,  4,  4, 40,  4, 40])