0

I have a pandas column of values with either 0 or 1, eg. 1, 1, 1, 0, 0, 1, 1, 0, ... I want to compute a new column that counts the consecutive 1 before it encounters 0, if it encounters 0, it will reset the counts.

data = {'input': [0, 1, 1, 1, 1, 0, 0, 1, 1],  'expected output': [0, 1, 2, 3, 4, 0, 0, 1, 2]}
df = pd.DataFrame.from_dict(data)
df[['input',  'expected output']]

enter image description here

# logic
lst_in = [0, 1, 1, 1, 1, 0, 0, 1, 1]
lst_out = []        
lst_out.append(lst_in[0])      # lst_out 1st element is the same as the 1st element of lst_in
x_last = lst_in[0]
y_last = 0
for x in lst_in[1:]:
    if x_last == 0:     # reset 
        y = x
        y_last = y

    elif x_last == 1:   # cum current y
        if x == 1:
            y = x + y_last
        elif x == 0:    # reset next 
            y = 0

    x_last = x
    y_last = y
    #print(x_last, y_last)
    lst_out.append(y)

print(lst_out)

I can make if work if I convert it to list first. However, I cannot figure out how to make the logic work under pandas framework

Haibin Chen
  • 113
  • 1
  • 8
  • Possible duplicate of [Python pandas cumsum with reset everytime there is a 0](https://stackoverflow.com/questions/45964740/python-pandas-cumsum-with-reset-everytime-there-is-a-0) – DjaouadNM Sep 02 '19 at 20:35

1 Answers1

1
df.groupby(df['input'].eq(0).cumsum()).input.cumsum()
Mark Wang
  • 2,623
  • 7
  • 15