如何用Numpy广播机制实现按行索引提取数组元素？

核心解决方案 #

你可以用Numpy的高级索引或take_along_axis函数直接实现，一行代码就能完成需求：

解法1：高级索引

output = NxM[np.arange(N)[:, None], Nxa]

解法2：take_along_axis函数（更直观）

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

完整代码验证 #

将解法代入你的示例代码，运行结果完全符合预期：

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]])

# 任选一种解法即可
output = NxM[np.arange(N)[:, None], Nxa]
# output = np.take_along_axis(NxM, Nxa, axis=1)

print(output)

运行输出：

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

原理说明 #

• 高级索引方式：np.arange(N)[:, None]生成形状为(N,1)的数组，和(N,a)的Nxa自动广播为(N,a)的索引对，每一组索引对应原数组的行位置和该行要提取的列位置，精准定位目标元素。
• take_along_axis方式：该函数专门用于沿指定轴按索引提取元素，axis=1指定沿列方向操作，Nxa中的每个值对应该行要提取的列索引，用法更直白易懂。

内容的提问来源于stack exchange，提问作者OliverBunting