我有一个2d的numpy数组:data.shape==(n,8),还有另一个ind.shape=(n,4)的数组。ind数组和data长度相等,包含像[4,3,0,6]这样的索引。我该如何创建另一个形状为(n,4)的数组,其中包含由ind指定的data元素?我的实际数组很长(shape[0]),所以循环速度很慢。必须有比循环更好的方法吧?
import numpy as np
# Example data
data = np.array([[ 0.44180102, -0.05941365, 2.1482739 , -0.56875081, -1.45400572,
-1.44391254, -0.33710766, -0.44214518],
[ 0.79506417, -2.46156966, -0.09929341, -1.07347179, 1.03986533,
-0.45745476, 0.58853107, -1.08565425],
[ 1.40348682, -1.43396403, 0.8267174 , -1.54812358, -1.05854445,
0.15789466, -0.0666025 , 0.29058816]])
ind = np.array([[3, 4, 1, 5],
[4, 7, 0, 1],
[5, 1, 3, 6]])
# This is the part I want to vectorize:
out = np.zeros(ind.shape)
for i in range(ind.shape[0]):
out[i,:] = data[i,ind[i,:]]
# This should be good
assert np.all(out == np.array([[-0.56875081, -1.45400572, -0.05941365, -1.44391254],
[ 1.03986533, -1.08565425, 0.79506417, -2.46156966],
[ 0.15789466, -1.43396403, -1.54812358, -0.0666025 ]]))