Подсчет последовательных 1 в массиве numpy

Question

[1, 1, 1, 0, 0, 0, 1, 1, 0, 0] 

У меня есть массив numpy, состоящий из 0 и 1, как указано выше. Как я могу добавить все последовательные 1, как показано ниже. Каждый раз, когда я сталкиваюсь с 0, я сбрасываю.

[1, 2, 3, 0, 0, 0, 1, 2, 0, 0] 

Я могу сделать это, используя цикл for, но есть ли более пифонические soln, используя numpy?

Answer 1

Вот Векторизованный подход -

def island_cumsum_vectorized(a): 
    a_ext = np.concatenate(([0], a, [0])) 
    idx = np.flatnonzero(a_ext[1:] != a_ext[:-1]) 
    a_ext[1:][idx[1::2]] = idx[::2] - idx[1::2] 
    return a_ext.cumsum()[1:-1] 

Пример запуск -

In [91]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0]) 

In [92]: island_cumsum_vectorized(a) 
Out[92]: array([1, 2, 3, 0, 0, 0, 1, 2, 0, 0]) 

In [93]: a = np.array([0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1]) 

In [94]: island_cumsum_vectorized(a) 
Out[94]: array([0, 1, 2, 3, 4, 0, 0, 0, 1, 2, 0, 0, 1]) 

тест Продолжительность

Для тайминги, я хотел бы использовать входной массив образца OP и повторите/плитка его и надеюсь, что это должно быть less opportunistic benchmark -

Малый корпус:

In [16]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0]) 

In [17]: a = np.tile(a,10) # Repeat OP's data 10 times 

# @Paul Panzer's solution 
In [18]: %timeit np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])]) 
10000 loops, best of 3: 73.4 µs per loop 

In [19]: %timeit island_cumsum_vectorized(a) 
100000 loops, best of 3: 8.65 µs per loop 

Bigger случай:

In [20]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0]) 

In [21]: a = np.tile(a,1000) # Repeat OP's data 1000 times 

# @Paul Panzer's solution 
In [22]: %timeit np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])]) 
100 loops, best of 3: 6.52 ms per loop 

In [23]: %timeit island_cumsum_vectorized(a) 
10000 loops, best of 3: 49.7 µs per loop 

Нет, я хочу действительно огромный случай:

In [24]: a = np.array([1, 1, 1, 0, 0, 0, 1, 1, 0, 0]) 

In [25]: a = np.tile(a,100000) # Repeat OP's data 100000 times 

# @Paul Panzer's solution 
In [26]: %timeit np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])]) 
1 loops, best of 3: 725 ms per loop 

In [27]: %timeit island_cumsum_vectorized(a) 
100 loops, best of 3: 7.28 ms per loop 

Answer 2

Если список понимание является приемлемым

np.concatenate([np.cumsum(c) if c[0] == 1 else c for c in np.split(a, 1 + np.where(np.diff(a))[0])])