数组的矢量化索引

Question

最初我有这样的东西：

a = 1   # Some randomly generated positive integer
b = -1  # Some randomly generated negative integer
c = 0   # Constant 0
i = 0  # Randomly picked from (0, 1, 2)
d = [a, b, c][i]

我想将其矢量化，以便可以生成许多样本

所以我有三个长度为N的数组，一个长度为N的索引数组，我想用这个索引数组从三个数组中选择一个

a = np.array([1, 2, 3, 4])
b = np.array([-1, -2, -3, -4])
c = np.array([0, 0, 0, 0])
i = np.array([2, 1, 2, 0])
d = np.array([a, b, c])[i]  # Doesn't work
# Would like the result:
d = np.array([0, -2, 0, 4])

d = a * (i == 0) + b * (i == 1) + c * (i == 2)是有效的，但肯定有一种方法看起来更像未经编码的

Answer 1

从三个数组中生成一个二维数组，然后使用整数索引

>>> e = np.vstack([a,b,c])
>>> i = np.array([2, 1, 2, 0])
>>> e[(i,np.arange(i.shape[0]))]
array([ 0, -2,  0,  4])
>>> 

Answer 2

注意你的答案在的对角线上

np.array([a, b, c])[i]

所以你可以去：

np.array([a, b, c])[i].diagonal()