如何使用索引数组为多维数组取值 替代逐元素遍历的for循环

有更简便且运行效率更高的实现方式，不需要手写for循环，可以直接利用数值计算库的高级索引特性实现，以下是不同使用场景的具体方案：

NumPy 数组场景 #

你可以任选以下两种写法：

• 写法1：配对行/列整数索引
  构造长度为32的行索引序列，和列索引I配对取值即可。由于你的I是形状为(32,1)的二维数组，需要先压缩为一维数组匹配行索引的维度：

  import numpy as np
  output = X[np.arange(32), I.squeeze()]
  

• 写法2：使用take_along_axis内置方法（维度兼容性更好，不易写错）
  不需要手动构造行索引，直接指定沿列维度（axis=1）用I取值即可：

  output = np.take_along_axis(X, I, axis=1)
  # 如果你需要和原for逻辑一致的、去掉最后一维的输出，再加.squeeze(axis=1)即可
  

PyTorch 张量场景 #

直接用内置的gather方法实现即可，逻辑和NumPy的take_along_axis一致：

import torch
output = X.gather(1, I)
# 同理需要降维的话加.squeeze(axis=1)即可

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