Python 嵌套使用numpy.where

我将数据存储在三个不同的数组中：A、B和C。
我还有一个整数数组Z，其中包含1、2或3。
A、 B、C和Z具有相同的形状。
我想创建一个新数组D，如果Z中的对应元素为1，则包含a中列出的值；如果Z中的对应元素为2，则包含B中列出的值；如果Z中的对应元素为3，则包含C中列出的值

为此，我使用嵌套版本的numpy.where编写了代码。然而，代码看起来很难看。有没有更好的方法可以做到这一点

import numpy as np
#Create arrays that hold the data
A = np.array([1.,1.,1.])
B = A * 2
C = A * 3
#Create an array that hold the zone numbers
Z = np.array([1,2,3])
#Now calculate the new array by figuring out the appropriate zone at each location
D = np.where(Z==1,A,np.where(Z==2,B, np.where(Z==3,C,0.0)))
#Output
print A
print B
print C
print D

[ 1.  1.  1.]
[ 2.  2.  2.]
[ 3.  3.  3.]
[ 1.  2.  3.]

看看这是否有效：

new=np.zeros_like(A)
new[Z==1] = A[Z == 1]
new[Z==2] = B[Z == 2]
new[Z==3] = C[Z == 3]

---

谢谢我想我会采用这种方法来提高代码的可读性。