嵌套使用numpy.where

Question

我有三个不同的数组，分别叫做 A、B 和 C。
还有一个整数数组 Z，里面的值只能是 1、2 或 3。
这四个数组 A、B、C 和 Z 的形状是一样的。
我想创建一个新的数组 D，如果 Z 中对应的位置是 1，就把 A 中的值放到 D 里；如果是 2，就把 B 中的值放到 D 里；如果是 3，就把 C 中的值放到 D 里。

为此，我写了一段代码，使用了嵌套的 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.]

Answer 1

看看这个能不能用：

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