python - NumPy - 捕捉前三个连续负数的索引

Question

我需要找到三个连续负数第一次出现的索引。在正常的 Python 方式中，我会这样做:

a = [1,-1,1,-1,1,-1,1,-1,-1,-1,1,-1,1]
b=0
for i,v in enumerate(a):
    if v<0:
        b+=1
    else:
        b=0
    if b==3:
        break
indx = i-2

有人知道如何以更智能的 NumPy 方式做到这一点吗？

Answer 1

这是一个借助卷积的矢量化解决方案 -

def first_consecutive_negative_island(a, N=3):
    mask = np.convolve(np.less(a,0),np.ones(N,dtype=int))>=N
    if mask.any():
        return mask.argmax() - N + 1
    else:
        return None

sample 运行-

In [168]: a = [1,-1,1,-1,1,-1,1,-1,-1,-1,1,-1,1]

In [169]: first_consecutive_negative_island(a, N=3)
Out[169]: 7

无论组在哪里都可以工作-

In [175]: a
Out[175]: [-1, -1, -1, 1, -1, 1, -1, -1, -1, 1, -1, 1]

In [176]: first_consecutive_negative_island(a, N=3)
Out[176]: 0

没有负数，它优雅地返回 None -

In [183]: a
Out[183]: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

In [184]: first_consecutive_negative_island(a, N=3)

对于恰好三个连续的负数搜索，我们可以使用切片，像这样-

def first_consecutive_negative_island_v2(a):
    m =  np.less(a,0)
    mask = m[:-2] & m[1:-1] & m[2:]
    if mask.any():
        return mask.argmax()
    else:
        return None

时间 -

In [270]: a = np.random.randint(-1,2,(1000000)).tolist()

In [271]: %timeit first_consecutive_negative_island(a, N=3)
10 loops, best of 3: 44.5 ms per loop

In [272]: %timeit first_consecutive_negative_island_v2(a)
10 loops, best of 3: 38.7 ms per loop