用另一个数组的行索引一个数组答案

【问题标题】：Index an array with the rows of another array用另一个数组的行索引一个数组
【发布时间】：2023-01-13 23:34:49
【问题描述】：

例如我有一个数组作为

a=    [[[98, 92, 92, 81],
        [98, 26, 98, 36],
        [83, 39, 96, 73],
        [75, 36, 37, 68],
        [56, 87, 62, 73]],

       [[98, 43, 36, 96],
        [92, 88, 94, 71],
        [ 7, 24, 33, 18],
        [14, 75, 30, 43],
        [91, 54, 74, 59]],

       [[60, 97, 52, 95],
        [78, 90, 99, 16],
        [59, 89,  5, 85],
        [44, 33, 52, 93],
        [41, 77, 64, 88]]]

和索引数组为

b= [[1, 4],
    [1, 3],
    [0, 4]]

然后我想要最终输出作为for循环的输出如下

for i in range(3):
        a[i, b[i]]

样本输出为

[[[98, 26, 98, 36],
  [56, 87, 62, 73]],
  
 [[14, 75, 30, 43],
  [92, 88, 94, 71]],

 [[41, 77, 64, 88],
  [60, 97, 52, 95]]]

我需要这个来将批量索引应用于批量数据。所以通常我会将一行索引向量应用于批处理数据中的一个矩阵（2D）

【问题讨论】：

底部的两帧交换了预期输出中的行。

标签： numpy-slicing

【解决方案1】：

对于您展示的案例，也许这样的事情会做：

import numpy as np
a = np.array(a)

使用索引数组b，选择元素：

parts  = [a[n,i,:] for n,i in enumerate(b)]

最后将零件列表转换为一个数组：

np.array(parts)

结果：

array([[[98, 26, 98, 36],
        [56, 87, 62, 73]],
       [[92, 88, 94, 71],
        [14, 75, 30, 43]],
       [[60, 97, 52, 95],
        [41, 77, 64, 88]]])

【讨论】：

【解决方案2】：

您可以使用直接索引来做到这一点。输出形状将是输入的广播形状加上带:的轴：

a[np.arange(3).reshape(-1, 1), b, :]

你也可以使用np.take_along_axis：

np.take_along_axis(a, b[..., None], axis=1)

【讨论】：