从另一个二维数组中提取带有索引的二维数组元素

Question

我有一个二维 numpy 数组：data.shape==(n,8) 和另一个 ind.shape=(n,4)。ind 数组的长度与数据相同，并包含 [4,3,0,6] 等索引。如何创建另一个包含由 ind 索引指定的数据元素的 shape==(n,4) 数组？我的实际数组很长（形状 [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 ]]))

Answer 1

如果我们索引到 raveleddata数组，这可以很容易地完成：

out = data.ravel()[ind.ravel() + np.repeat(range(0, 8*ind.shape[0], 8), ind.shape[1])].reshape(ind.shape)

解释

如果将其分解为三个步骤可能更容易理解：

indices = ind.ravel() + np.repeat(range(0, 8*ind.shape[0], 8), ind.shape[1])
out = data.ravel()[indices]
out = out.reshape(ind.shape)

ind有data我们想要的元素的信息。不幸的是，它以二维索引表示。上面的第一行将这些转换indices为 1-D raveled data。上面的第二行从 raveled array 中选择那些元素data。第三行将二维形状恢复为out。由 表示的二维索引ind转换为 indindices具有索引