Pandas - How to assign different classes to a recurring series without using a loop?

Viewed 41

I have a series with recurring values in pandas dataframe. The first 20 values are given below:

# input
[-1,  1, -1,  1,  1, -1,  1, -1,  1, -1,  1,  1,  1,  1,  1,  1, -1, 1,  1,  1]

Note that -1 indicates the start of a new class while 1 indicates continuation of previous class. I can change "-1" and "1" to any value for the encoding. What I want for the output is to be class labels like given below:

# output
[0, 0, 1, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5]

For now I have done this using a for loop. I could not find a better way with in-built pandas / python functions.

How do I achieve this without using for loop?

1 Answers

Compare values by -1 and use cumulative sum by Series.cumsum and if first value of Series is 0 subtract 1:

s = pd.Series([-1,  1, -1,  1,  1, -1,  1, -1,  1, -1,  1,  1, 
                1,  1,  1,  1, -1, 1,  1,  1])

s1 = s.eq(-1).cumsum()
out = s1.sub(1) if s[0] == -1 else s1

If always first value is -1 subtract 1:

out = s.eq(-1).cumsum().sub(1)
Related