从指定索引数组的多维数组中提取元素的最快方法

Question

假设我有一个 n 维数组 A，其中有一个任意 n：

a = np.arange(3*4*5*...).reshape((3, 4, 5, ...))

...而且我有一个索引数组（正是这个形状或转置的）：

ix = np.array([
  [0,1,1,...],
  [0,0,1,...],
  [1,0,1,...],
])

从 ix 的 at 索引中获取元素数组的最快方法是什么？

（例如最快的获取方式：

def take_at(a, ix):
  res = np.empty((ix.shape[0],*a.shape[ix.shape[1]:]))
  for i in range(ix.shape[0]):
    res[i,...] = a[(*ix[i],)]
  return res

...仅使用（如果可能的话）numpy 向量化函数/循环/操作？）

用于测试的简单复制粘贴示例：：

a = np.arange(3*4*5).reshape((3,4,5))

ix = np.array([
  [0,1,1],
  [0,0,1],
  [1,0,1],
])

def take_at(a, ix):
  res = np.empty((ix.shape[0],*a.shape[ix.shape[1]:]), dtype=a.dtype)
  for i in range(ix.shape[0]):
    res[i,...] = a[(*ix[i],)]
  return res

take_at(a, ix)

Answer 1

你需要将ix打包成一个元组：

take_at(a, ix)
Out[]: array([ 6,  1, 21])

a[tuple(ix.T)]
Out[]: array([ 6,  1, 21])

Answer 2

你可以这样做：

a[ix[:,0],ix[:,1],ix[:,2]]

输出：

array([ 6,  1, 21])