获取锯齿状切片的 numpy 简写

Question

我有一个经常做的操作，我称之为“锯齿状切片”，因为我不知道它的真实名称。最好通过例子来解释：

a = np.random.randn(50, 10)
entries_of_interest = np.random.randint(10, size = 50)  # Vector of 50 indices between 0 and 9
# Now I want the values contained in each row of a at the corresponding index in "entries of interest"
jagged_slice_of_a = a[np.arange(a.shape[0]), entries_of_interest]
# jagged_slice_of_a is now a vector with 50 elements.  Good.

唯一的问题是执行此a[np.arange(a.shape[0]), entries_of_interest]索引有点麻烦（为此而必须构造“np.arange(a.shape[0])”似乎很愚蠢）。我想要类似:操作员的东西，但它:会做其他事情。有没有更简洁的方法来执行此操作？

最佳答案：

不，原生 numpy 没有更好的方法。如果需要，您可以为此创建一个辅助函数。

Answer 1

这很麻烦，只是因为它需要更多的输入来完成对您来说似乎很简单的任务。

a[np.arange(a.shape[0]), entries_of_interest]

但正如您所注意到的，语法上更简单的a[:, entries_of_interest]在 中还有另一种解释numpy。选择数组列的子集是比从每行选择一个（随机）项目更常见的任务。

你的案例只是一个特殊的例子

a[I, J]

其中I和J是两个形状相同的数组。在一般情况下entries_of_interest，可能小于a.shape[0]（不是所有行），或大于（某些行中的多个项目），甚至是 2d。它甚至可以重复选择某些元素。

我在其他 SO 问题中发现，当应用于a.flat. 但这需要一些数学来构建这种I*n+J扁平指数。

凭借您对 的专业知识J，构建I似乎是额外的工作，但numpy不能做出这种假设。如果这种选择更常见，有人可以编写一个包装您的表达式的函数

def  peter_selection(a,I):
   # check the a.shape[0]==I.shape[0]
   return a[np.arange(a.shape[0]), I]