python - 用于锯齿状切片的 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]