如何用 2 个索引数组索引高维数组

Question

我有一个 (2 x 1 x 2 x 2 x 2) 维数组：

array([[[[[ 7.,  9.],
          [10., 11.]],

         [[19., 18.],
          [20., 16.]]]],



       [[[[24.,  5.],
          [ 6., 10.]],

         [[18., 11.],
          [45., 12.]]]]])

最后两个维度分别是H（高度）和W（宽度）。现在我有两个单独的数组，索引沿 H 和 W：

idx2=np.array([1, 1, 0, 1]) # index along H
idx3=np.array([1, 0, 0, 0]) # index along W

因此，就最后两个维度而言，我想从[[ 7.,9.],[10.,11.]]中提取第(1,1)个元素，即11；以及来自[[19.,18.],[20.,16.]] 的第 (1,0) 个元素，即 20，依此类推。最终结果应该是一个(2 x 1 x 2)数组：

array([[[11., 20.]],

       [[24., 45.]]])

感谢您的帮助！

Answer 1

查看numpy indexing：

idx0=np.array([0,0,1,1])
idx1=np.array([0,1,0,1])
array[idx0,0,idx1,idx2,idx3]

另一种更 hacky 的方式：（我在我的代码中将数组命名为 a）

np.diag(a[:,:,:,idx2,idx3].reshape((4,-1))) 

在一般情况下，您必须将 4 替换为其他维度的乘积。

Answer 2

使用np.ravel_multi_index 和np.take_along_axis 的组合解决此问题的可能方法。

---

您可以解开最后两个维度上的索引，并在此扁平空间维度 h x w 轴上使用 np.take_along_axis：

>>> flat_idx = (idx2*x.shape[-1]+idx3).reshape(*x.shape[:-2], 1)
>>> flat_idx 
array([[[[3],
         [2]]],


       [[[0],
         [2]]]])

或者，你可以选择使用内置的np.ravel_multi_index，但是会稍微长一些：

>>> flat_idx = np.ravel_multi_index((idx2, idx3), x.shape[-2:]).reshape(*x.shape[:-2], 1)

然后展平x 的最后两个维度并收集索引：

>>> res = np.take_along_axis(x.reshape(*x.shape[:-2], -1), flat_idx, -1)
>>> res
array([[[[11.],
         [20.]]],


       [[[24.],
         [45.]]]])

此时需要重塑：

>>> res.reshape(*x.shape[0:-2])
array([[[11., 20.]],

       [[24., 45.]]])

---

如果你推断形状，idx2 和 idx3 这归结为：

>>> flat_idx = (idx2*2+idx3).reshape(2, 1, 2, 1)
>>> res = np.take_along_axis(x.reshape(2, 1, 2, 4), flat_idx, -1)
>>> res.reshape((2, 1, 2))

---

上述方法可用于处理idx2、idx3、idx4、...等更一般的情况

Answer 3

如果你有一个 2x1x2x2x2 维数组，array_name。

import numpy as np
array_name = np.array([[[[[ 7.,  9.],
                          [10., 11.]],
                         [[19., 18.],
                          [20., 16.]]]],
                       [[[[24.,  5.],
                          [ 6., 10.]],
                         [[18., 11.],
                          [45., 12.]]]]])

array_name[0] 和 array_name[1] 都是 1x2x2x2 维数组。

array_name[0].shape # (1, 2, 2, 2)

array_name[0][0] 是一个 2x2x2 维数组。

array_name[0][0].shape # (2, 2, 2)

array_name[0][0][0] 是一个 2x2 维数组。

array_name[0][0][0].shape # (2, 2)

array_name[0][0][0][0] 是一个二维数组。

array_name[0][0][0][0].shape # (2,)

对于问题中给出的数组：

array_name[0][0][0][1][1] # 11.
array_name[0][0][1][1][0] # 20.
array_name[1][0][0][0][0] # 24.
array_name[1][0][1][1][0] # 45.

此方法非常冗长，因此您可以使用高级索引。

idx_0 = [0,0,1,1]
idx_2 = [0,1,0,1]
idx_3 = [1,1,0,1]
idx_4 = [1,0,0,0]

每个idx_X 数组中的值取自上述详细代码中的数组索引，并用于创建一个包含 4 个元素的数组，如下所示。

array_name[idx_0, 0, idx_2, idx_3, idx_4]
# array([11., 20., 24., 45.])