是否有一种简单的 numpy 方法来广播一个维度并使用其他维度进行查找？答案

【问题标题】：Is there a simple numpy way to broadcast one dimension and look-up using others?是否有一种简单的 numpy 方法来广播一个维度并使用其他维度进行查找？
【发布时间】：2023-01-13 23:06:46
【问题描述】：

我有 2 个尺寸为 NxM 和 Nxa 的 numpy 数组。我想从 NxM 矩阵的每一行中提取“M”个元素中的“a”个，行索引由 Nxa 给出。似乎应该有一个简单的广播解决方案，但我找不到。

希望下面给出的示例具有示范性。

import numpy as np
N=5
M=7
a=3

NxM = np.array([[0, 1, 2, 3, 4, 5, 6],
                [10, 11, 12, 13, 14, 15, 16],
                [20, 21, 22, 23, 24, 25, 26],
                [30, 31, 32, 33, 34, 35, 36],
                [40, 41, 42, 43, 44, 45, 46]])

Nxa = np.array([[1, 3, 6],
                [0, 1, 2],
                [1, 4, 5],
                [2, 3, 4],
                [3, 5, 6]])

#desired output
output = [[1, 3, 6],
          [10, 11, 12],
          [21, 24, 25],
          [32, 33, 34],
          [43, 45, 46]]

【问题讨论】：

NxM[np.arange(NxM.shape[0])[:,None], Nxa]

标签： python numpy array-broadcasting

【解决方案1】：

您正在寻找花哨的索引：

NxM[np.arange(N).reshape(-1, 1), Nxa]

输出数组的形状是广播索引的形状。

【讨论】：

【解决方案2】：

numpy.take_along_axis 就是您要找的。

np.take_along_axis(NxM, Nxa, axis=1)

array([[ 1,  3,  6],
       [10, 11, 12],
       [21, 24, 25],
       [32, 33, 34],
       [43, 45, 46]])

【讨论】：