NumPy高级索引实现：按指定索引提取数组元素

嘿，你的这个问题我太熟啦！先给你点个赞，你想到的reshape解法完全是可行的，但确实还有更简洁、可读性更强的实现方式，我给你梳理几种：

• 你的现有解法：完全没问题！
  你写的这段代码逻辑清晰，通过A.reshape(-1, K)把三维数组拉成二维，再用扁平化的index数组提取对应位置的元素，最后reshape回(N,M)形状，计算效率也不错，适合习惯扁平化处理的场景。

  B = A.reshape(-1, K)[np.arange(N * M), index.reshape(-1)].reshape(N, M)
  

• 最简洁直观的实现：NumPy高级索引（花式索引）
  这是我最推荐的写法，直接利用NumPy的多维索引特性，生成对应前两个维度的坐标数组，配合index数组直接提取元素，完全不需要reshape，可读性拉满：

  import numpy as np
  # 假设已定义A（形状(N,M,K)）和index（形状(N,M)）
  i, j = np.indices((N, M))
  B = A[i, j, index]
  

  这里np.indices((N,M))会生成两个形状为(N,M)的数组，分别对应每个位置的行、列索引，和index数组一起作为A的三维索引，直接就能得到符合要求的(N,M)形状的B数组。

• 另一种紧凑的写法：手动广播生成索引
  如果不想用np.indices，也可以手动生成索引并利用广播机制，代码更紧凑：

  i = np.arange(N)[:, None]  # 将一维数组转为(N,1)的列向量，实现广播
  B = A[i, np.arange(M), index]
  

  这里np.arange(N)[:, None]和np.arange(M)会自动广播成(N,M)形状的i、j索引数组，效果和上面的方法完全一致。

• 几种方法对比：

  • 你的reshape解法：效率不错，但可读性稍弱，适合熟悉数组扁平化操作的开发者。
  • 高级索引解法：语义最清晰，代码一看就懂，而且NumPy内部对高级索引做了优化，效率和reshape方法不相上下，是首选方案。

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