Łączne zliczenia w NumPy bez iteracji

Question

mam tablicą tak:

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

Chciałbym przeciwdziałają obrotowej przypadkach 1,0 że resetuje gdy napotka 0,0, tak wynik byłby:

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

Moją pierwszą myślą było użyć czegoś podobnego b = np.cumsum (a [a == 1,0]), ale nie wiem jak (1) modyfikować to, aby przywrócić w zera lub (2) jak to zorganizować, więc ou tablica tput ma taki sam kształt jak tablica wejściowa. Wszelkie pomysły, jak to zrobić bez iteracji?

Answer 1

Myślę, że można zrobić coś takiego

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 

który daje mi

>>> 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]) 

To działa, ponieważ możemy użyć skumulowany maksimum, aby dowiedzieć się "overcount":

>>> 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]) 

---

Poczytalność Testy:

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!") 

a następnie

>>> test() 
hooray!