python - numpy 逻辑中首次出现

Question

假设我有一个 numpy.ndarray:

a = np.array([0,4,10,0,11,10])

我将其与 10 进行了比较。

a >= 10
# array([False, False,  True, False,  True,  True], dtype=bool)

我想要一个 True，即仅在第一次出现时为 True。

我想将其应用于n-D numpy.ndarray 中的给定轴。(例如，1000*1000*10)

a_2d = np.array([[0,4,10],[0,11,10]])
#if axis == 1: array([[False, False, True], [False, True, False]])

我做了什么:

对于一维数组，我设法使用 this 来做到这一点.

b=np.zeros(a.size)
b[np.argmax(a>=10)]=True
#b=array([ 0.,  0.,  1.,  0.,  0.,  0.])

但是，我不知道如何将其应用于大型n维数组。

Answer 1

对于一维或二维，这个应该不需要 for 循环:

def firstByRow(a, f = lambda x: x >= 10):
    b = (np.cumsum(f(a), axis = -1) == 1).T
    b[1:] = b[1:] * np.equal(b[1:], np.diff((f(a)).astype(int), axis = -1).T)
    return b.T

不确定它是否会比稍微循环的代码更快，因为它同时具有 cumsum 和 diff

编辑:

您也可以这样做，这可能会更快(利用 np.unique(return_index = True) 选择第一个匹配项):

def firstByAxis(a, f = lambda x: x >= 10, axis = 0):
    c = np.where(f(a))
    i = np.unique(c[axis], return_index = True)[1]
    b = np.zeros_like(a)
    b[tuple(np.take(c, i, axis = -1))] = 1
    return b