NumPy 中切片的不同起始索引

【问题标题】：Different starting indices for slices in NumPyNumPy 中切片的不同起始索引
【发布时间】：2016-04-13 20:10:48
【问题描述】：

我想知道是否可以不使用 for 循环进行迭代来执行以下操作：

a = np.array([[1, 2, 5, 3, 4],
              [4, 5, 6, 7, 8]])

cleaver = np.argmax(a == 5, axis=1)  # np.array([2, 1])

foo(a, cleaver)

>>> np.array([False, False, True, True, True],
             [False, True, True, True, True])

有没有办法通过切片或其他非迭代函数来实现这一点？我使用的数组非常大，逐行迭代它们非常昂贵。

【问题讨论】：

标签： python arrays performance numpy vectorization

【解决方案1】：

你可以使用一些broadcasting魔法 -

cleaver[:,None] <= np.arange(a.shape[1])

示例运行 -

In [60]: a
Out[60]: 
array([[1, 2, 5, 3, 4],
       [4, 5, 6, 7, 8]])

In [61]: cleaver
Out[61]: array([2, 1])

In [62]: cleaver[:,None] <= np.arange(a.shape[1])
Out[62]: 
array([[False, False,  True,  True,  True],
       [False,  True,  True,  True,  True]], dtype=bool)

【讨论】：