Recuentos acumulados en NumPy sin iteración

Tengo una matriz como tal:

a = np.array([0.1, 0.2, 1.0, 1.0, 1.0, 0.9, 0.6, 1.0, 0.0, 1.0]) 

Me gustaría tener un contador de ejecuciones de 1.0 que se reinicie cuando encuentre un 0.0 , por lo que el resultado sería:

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

Mi idea inicial fue usar algo como b = np.cumsum (a [a == 1.0]), pero no sé cómo (1) modificar esto para restablecer en ceros o (2) exactamente cómo estructurarlo la matriz de salida tiene la misma forma que la matriz de entrada. ¿Alguna idea de cómo hacer esto sin iteración?

Creo que podrías hacer algo como

 def rcount(a): without_reset = (a == 1).cumsum() reset_at = (a == 0) overcount = np.maximum.accumulate(without_reset * reset_at) result = without_reset - overcount return result 

lo que me da

 >>> a = np.array([0.1, 0.2, 1.0, 1.0, 1.0, 0.9, 0.6, 1.0, 0.0, 1.0]) >>> rcount(a) array([0, 0, 1, 2, 3, 3, 3, 4, 0, 1]) 

Esto funciona porque podemos usar el máximo acumulado para calcular el “recuento excesivo”:

 >>> without_reset * reset_at array([0, 0, 0, 0, 0, 0, 0, 0, 4, 0]) >>> np.maximum.accumulate(without_reset * reset_at) array([0, 0, 0, 0, 0, 0, 0, 0, 4, 4]) 

Pruebas de cordura:

 def manual(arr): out = [] count = 0 for x in arr: if x == 1: count += 1 if x == 0: count = 0 out.append(count) return out def test(): for w in [1, 2, 10, 10**4]: for trial in range(100): for vals in [0,1],[0,1,2]: b = np.random.choice(vals, size=w) assert (rcount(b) == manual(b)).all() print("hooray!") 

y entonces

 >>> test() hooray!